A Tensor Is a Flat Buffer and a Promise

How PyTorch maps a logical index like x[b, h, s, d] to an address in flat memory — and why views, reshapes, and packed layouts can hand you the right shape with the wrong values. The one idea that turns tensor gymnastics from memorized incantations into arithmetic you can derive.
fundamentals
tensors
Author

Maitreya Suin

Published

July 14, 2026

Here are eight numbers. They came out of a transformer’s fused attention projection, packed into one tensor, and your job — the job of the very next line of code — is to split them back into the two queries and two keys they represent.

import torch
packed = torch.tensor([10., 11., 20., 21., 30., 31., 40., 41.])

The layer that produced this used two attention heads, and each query and key is a 2-dimensional vector. So there are four little vectors hiding in those eight numbers. But which four? Here are two completely reasonable stories:

flowchart TB
    T["<b>Same eight values, two packing contracts</b><br/>The address contract is identical; the semantic interpretation is not."]

    A["<b>Story A — interleaved by head</b><br/>[ Q0: 10,11 | K0: 20,21 | Q1: 30,31 | K1: 40,41 ]<br/>Decode: view(..., n_heads, 2, d_head)"]
    B["<b>Story B — blocked by role</b><br/>[ Q0: 10,11 | Q1: 20,21 | K0: 30,31 | K1: 40,41 ]<br/>Decode: chunk(2, -1), then reshape each half"]

    C["Both decoders return the same shape.<br/><b>K0 is [20,21] in Story A but [30,31] in Story B."]
    
    T --> A
    T --> B
    A --> C
    B --> C

    classDef title fill:#f8fafc,stroke:#334155,color:#0f172a;
    classDef packing fill:#eff6ff,stroke:#2563eb,color:#1e3a8a;
    classDef warning fill:#fef2f2,stroke:#dc2626,color:#7f1d1d;
    classDef audit fill:#f0fdf4,stroke:#16a34a,color:#14532d;
    class T title;
    class A,B packing;
    class C warning;
    class D audit;

Look at K_head0 under each story. Story A says it’s [20, 21]. Story B says it’s [30, 31]. These are different vectors. They will produce different attention weights, a different output, a different trained model. And here is the uncomfortable part: whichever story you assume, the code you write to split the tensor will run without error, and will hand back a result of exactly the shape you expected. Nothing crashes. Nothing warns. If you guessed wrong, your model is quietly, permanently a little bit broken, and you will spend a weekend blaming your learning rate.

By the end of this post you’ll be able to write the split that resolves this correctly, prove it’s correct against a reference that can’t lie, and — more importantly — you’ll understand why the wrong version was still shaped right. That last part is the real subject. It rests on one idea, which I want to state up front because everything else hangs off it:

A tensor carries two contracts, not one. The address contract — shape, strides, offset — tells you where each element lives in memory. The packing contract — a convention the producing code chose and rarely writes down — tells you what each position means. view and reshape and transpose are fluent in the first contract and completely ignorant of the second. Correct tensor code needs both, and only one of them is checkable by the machine.

Most tensor bugs, and most tensor interview questions, live in the gap between those two contracts. Let’s build up the machinery to see the gap clearly.

TL;DR — A dense tensor is a 1-D block of memory plus (shape, strides, offset). Operations like view, transpose, permute, expand, and slicing don’t move any data; they hand you new metadata pointing at the same buffer. From that single fact you can derive why a flattening view often fails after a nontrivial transpose, why reshape can silently cost you a full memory copy, and why the “same shape, wrong values” bug above is not only possible but common.

The address contract: where a number lives

Forget nested lists. That’s how tensors print, but it is not how they exist. A dense tensor is a flat, one-dimensional run of memory, and a small bundle of metadata that explains how to read a multi-dimensional index off that flat run. Three pieces:

  • shape — the size of each dimension, e.g. (2, 3, 4).
  • strides — for each dimension, how many elements you step forward in the flat buffer to advance that index by one.
  • storage offset — how far into the buffer this tensor’s first element sits (usually 0, but not after a slice).

The rule connecting a logical index to a physical location is a single affine formula. For an index \((i_0, i_1, \dots)\) with strides \((s_0, s_1, \dots)\):

\[\text{address} = \text{offset} + \sum_k i_k \, s_k\]

That formula is the entire address contract. There’s nothing else to it — no per-element bookkeeping, no nested pointers. Which means you can compute an element’s location by hand and check yourself:

x = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4)

print("shape  ", tuple(x.shape))
print("strides", x.stride())          # (12, 4, 1)
print("offset ", x.storage_offset())

i, j, k = 1, 2, 3
addr = x.storage_offset() + i*x.stride(0) + j*x.stride(1) + k*x.stride(2)
print(f"\nx[{i},{j},{k}]        = {x[i,j,k].item()}")
print(f"address from formula = {addr}")
shape   (2, 3, 4)
strides (12, 4, 1)
offset  0

x[1,2,3]        = 23.0
address from formula = 23
shape   (2, 3, 4)
strides (12, 4, 1)
offset  0

x[1,2,3]        = 23.0
address from formula = 23

flowchart LR
    A["Logical coordinate x[1, 2, 3]"] --> B["0 + 1×12 + 2×4 + 3×1 = 23"]
    B --> C["Flat storage slot 23"]

The value at logical position [1,2,3] is 23.0, and our hand-computed address is 23. Because this tensor happens to be laid out in plain row-major order starting at offset 0, its logical order and its storage order coincide, and address 23 is genuinely the 24th slot in the buffer. (That coincidence is exactly what stops being true the moment we transpose, in a second — so don’t lean on it as a way to peek at raw memory in general. It works here because x is contiguous and offset-zero.)

Read the strides (12, 4, 1) against the shape (2, 3, 4) and they tell a physical story. The last axis has stride 1: adjacent elements along it are adjacent in memory. The middle axis has stride 4: stepping it jumps over one full row of 4 elements. The first axis has stride 12: stepping it jumps a whole 3×4 slab. So the last dimension is the fastest-varying — walk it and you walk memory in order. This is what “C-contiguous” (row-major) means, and it’s the layout every freshly built tensor gets.

Two footnotes before we move, because I told you the model is exactly this formula and I want to keep that honest. Strides are counted in elements, not bytes — a detail that trips up anyone coming from NumPy, where .strides is reported in bytes (a float64 array has strides like (24, 8) where the corresponding torch tensor reads (3, 1)). And this affine-address picture is the model for dense strided tensors, which is essentially everything you’ll handle; sparse, quantized, and nested tensors carry different metadata and don’t obey this formula. With those two caveats noted, the formula is the whole address contract.

Views: new metadata, same memory

Here’s where the address contract starts paying off. Because a tensor is just a buffer plus that metadata, an enormous number of operations can be performed without touching the buffer at all — they simply hand you a new bundle of metadata pointing at the same memory. These are called views, and transpose is the cleanest example. Transposing two axes doesn’t rearrange any numbers; it swaps their two sizes and their two strides, and that’s the entire operation:

y = x.transpose(1, 2)                 # swap axes 1 and 2: shape (2, 3, 4) -> (2, 4, 3)

print("x.stride() ", x.stride())      # (12, 4, 1)
print("y.stride() ", y.stride())      # (12, 1, 4)  <- the 4 and the 1 changed places
print("y is contiguous?", y.is_contiguous())
x.stride()  (12, 4, 1)
y.stride()  (12, 1, 4)
y is contiguous? False
x.stride()  (12, 4, 1)
y.stride()  (12, 1, 4)
y is contiguous? False

Same buffer, new lens. No copy, no arithmetic — just a relabeling of how to index the bytes that were already there. This is why transpose and permute are effectively instantaneous even on enormous tensors: their cost doesn’t scale with the data, because they don’t read the data.

But look at what it did to the strides: (12, 1, 4). They no longer decrease left to right. Walking the last axis of y now steps 4 elements at a time through memory, hopping around instead of gliding. For y.shape == (2, 4, 3), an ordinary row-major contiguous layout would have strides (12, 3, 1). Instead, y has strides (12, 1, 4): advancing the final logical axis jumps four storage elements rather than one. Therefore y is not contiguous in torch.contiguous_format. y is a perfectly valid tensor; it just carries a memory layout that some operations can’t accept as-is, which is the crux of the next section.

Confirming they really share memory turns out to be slightly subtler than it looks, and the subtlety is instructive. The obvious check — compare the address of the first element — can disagree with the truth, because two views of the same buffer can start at different offsets:

def shares_storage(a, b):
    """Do these two tensors point into the same underlying allocation?"""
    return a.untyped_storage().data_ptr() == b.untyped_storage().data_ptr()

base = torch.arange(12).reshape(3, 4)
crop = base[1:, 1:]                    # a slice: a view starting partway in

print("first-element addresses equal? ", base.data_ptr() == crop.data_ptr())
print("same underlying storage?       ", shares_storage(base, crop))
print("crop's storage offset:", crop.storage_offset(), "| crop's strides:", crop.stride())
first-element addresses equal?  False
same underlying storage?        True
crop's storage offset: 5 | crop's strides: (4, 1)
first-element addresses equal?  False
same underlying storage?        True
crop's storage offset: 5 | crop's strides: (4, 1)

crop is a view into base — no data was copied — yet base.data_ptr() == crop.data_ptr() says False, because crop’s first element sits at offset 5 (skip one row of 4, then one more column) and so has a different element address even though it’s the same allocation. The honest question is “same storage?”, and untyped_storage().data_ptr() answers it. data_ptr() alone answers a narrower question — “same starting point?” — which is why I’ll use the shares_storage helper for the rest of this post. (NumPy folds this into one function, np.shares_memory(a, b), which is worth knowing if you switch between the two libraries.)

That crop example also quietly demonstrates the offset in action: slicing is a view, and a nonzero storage offset is how a view “starts partway into” a buffer without copying.

When view refuses and reshape pays

Now the practical consequence that every PyTorch user meets, usually as an error message. There are two ways to ask for a new shape, and they differ in exactly one respect: what they do when the layout won’t cooperate.

view demands a new shape that the current strides can already express, and hands you a view. If the memory isn’t laid out compatibly, it refuses — it will not silently rescue you:

y.view(-1)      # y is the non-contiguous transpose from earlier
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[5], line 1
----> 1 y.view(-1)      # y is the non-contiguous transpose from earlier

RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.
RuntimeError: view size is not compatible with input tensor's size and stride
(at least one dimension spans across two contiguous subspaces).
Use .reshape(...) instead.

The error is telling you something precise: to flatten y into one axis, memory would have to be traversed in an order its strides can’t produce in a single linear sweep, so no valid view exists. view won’t manufacture one.

reshape asks for the same thing, but with a fallback: if a view is possible it returns one, and if not, it quietly makes a contiguous copy and reshapes that. It succeeds where view failed — but “succeeds” now potentially means “allocated a whole new buffer and copied every element into it”:

z = y.reshape(-1)
print("reshape succeeded. does z share y's storage?", shares_storage(z, y))
reshape succeeded. does z share y's storage? False
reshape succeeded. does z share y's storage? False

False — this particular reshape copied. That’s not a bug; flattening a transposed tensor genuinely requires reordering memory, and reshape did the honest thing. But it did it silently, and in a training loop that silence is a cost you might not have budgeted for.

So the rule isn’t “use reshape, it’s more forgiving.” The rule is about intent:

  • Reach for view when a copy would be a bug — inside a hot loop, or when you’re about to write into the result and need it to alias the original. view failing is then a useful alarm.
  • Reach for reshape when you just want the shape and don’t care how you get it. But don’t assume it returned a view; if aliasing matters, check with shares_storage, because whether any given reshape copied depends on the input’s layout and is not something you should read off the syntax. (PyTorch’s own docs make this explicit: don’t depend on whether reshape aliased.)

One precise footnote on when view can succeed, for the curious: collapsing axes \(i\) and \(i{+}1\) into one is layout-compatible exactly when \(\text{stride}[i] = \text{stride}[i{+}1] \times \text{size}[i{+}1]\) — i.e. when the two axes are already adjacent in memory with no gap. That’s the whole condition. Contiguous tensors satisfy it for every adjacent pair, which is why view “just works” on them and starts refusing the moment a transpose or slice breaks the pattern. (NumPy exposes the same distinction differently: np.reshape mirrors torch.reshape, but assigning to arr.shape directly is a view-like operation that raises rather than copies — the equivalent of forcing the view path.)

Stride zero: expand versus repeat

One more view trick, because it’s the sharpest illustration that strides can lie in a useful way, and because it’s the direct setup for a real bug in a later post. Suppose you have a column vector of shape (3, 1) and you want to broadcast it out to (3, 4) — the same three values, repeated across four columns. The naive mental model is “make four copies.” The efficient reality is “set a stride of zero”:

a = torch.arange(3.).view(3, 1)
e = a.expand(3, 4)      # a broadcast VIEW
r = a.repeat(1, 4)      # an actual COPY

print("a.expand(3,4).stride()", e.stride(), " <- stride 0 on the stretched axis")
print("a.repeat(1,4).stride()", r.stride())
print("expand shares a's storage?", shares_storage(e, a))
print("repeat shares a's storage?", shares_storage(r, a))
print("same values though?       ", torch.equal(e, r))
a.expand(3,4).stride() (1, 0)  <- stride 0 on the stretched axis
a.repeat(1,4).stride() (4, 1)
expand shares a's storage? True
repeat shares a's storage? False
same values though?        True
a.expand(3,4).stride() (1, 0)  <- stride 0 on the stretched axis
a.repeat(1,4).stride() (4, 1)
expand shares a's storage? True
repeat shares a's storage? False
same values though?        True

A stride of 0 means “stepping along this axis doesn’t move in memory at all” — so all four columns read the very same three elements. expand creates the broadcasted view without allocating storage for repeated elements, repeat physically writes out all twelve values. They contain identical numbers, so for any read-only use — feeding into a matmul, a comparison, an addition — expand is usually the right starting point.

The catch, and the reason this matters later: because expand’s columns are the same memory, writing to them does something you almost never intend. This is subtler than “it raises an error” — for a single indexed write, PyTorch doesn’t stop you at all; it writes through the shared stride-0 axis and quietly changes every aliased element at once:

e[0, 1] = 99.       # try to set ONE element...
print("row 0 after e[0,1] = 99. :", e[0].tolist())
row 0 after e[0,1] = 99. : [99.0, 99.0, 99.0, 99.0]
row 0 after e[0,1] = 99. : [99.0, 99.0, 99.0, 99.0]

flowchart LR
    S0["storage a[0]"]
    S1["storage a[1]"]
    S2["storage a[2]"]

    subgraph E["Expanded view e"]
      direction TB
      R0["e[0,0]   e[0,1]   e[0,2]   e[0,3]"]
      R1["e[1,0]   e[1,1]   e[1,2]   e[1,3]"]
      R2["e[2,0]   e[2,1]   e[2,2]   e[2,3]"]
    end
    style E fill:#e2f0d9,stroke:#70ad47,stroke-width:1px

    W["Write e[0,1] = 99"]
    O["Logical row 0 becomes<br/><b>[99, 99, 99, 99]</b>"]

    S0 -->|all four cells alias this element| R0
    S1 -->|all four cells alias this element| R1
    S2 -->|all four cells alias this element| R2
    R0 --> W --> O

    classDef storage fill:#e0f2fe,stroke:#0284c7,color:#0c4a6e;
    classDef view fill:#f8fafc,stroke:#475569,color:#0f172a;
    classDef write fill:#fff7ed,stroke:#f97316,color:#7c2d12;
    classDef result fill:#fee2e2,stroke:#dc2626,color:#7f1d1d;
    class S0,S1,S2 storage;
    class R0,R1,R2 view;
    class W write;
    class O result;

All four columns of row 0 became 99, because they were never four elements — they were one element read four times, and you just overwrote it. No error, total corruption of the row. PyTorch does refuse the ops it can detect as unsafe up front — a whole-tensor in-place like e.add_(1) raises RuntimeError: unsupported operation: more than one element of the written-to tensor refers to a single memory location — but the single-element write above slips through and is exactly the kind of silent aliasing bug that ruins an afternoon.

So: expand to read a value in many places; repeat (or clone) the moment the copies must become independently writable. And treat “free” as applying to reads only — a downstream kernel may still copy an expanded tensor to get a contiguous buffer, so even the read case isn’t unconditionally free, just usually.

We now have the whole address contract in hand: strides place elements, views relabel without copying, view/reshape differ on whether they’ll pay for a copy, and stride-0 broadcasting. Everything so far has been about where numbers live. Time for the second contract — what they mean — and the bug from the opening.

The packing contract: what a position means

Return to those eight numbers. The address contract can tell you where the 5th element is stored. It cannot tell you whether the 5th element is part of a query or a key, because that isn’t a fact about memory — it’s a convention the producing code chose. Whoever wrote the attention layer decided how to interleave the heads and roles when they flattened everything into one axis. That decision is the packing contract, and it is almost never written down anywhere you can see it; you infer it from the code that produced the tensor, or from a paper, or from a comment if you’re lucky.

Let me make the two contracts concrete and separate, because conflating them is the whole trap:

  • Address / physical layout — strides, offset, contiguity. Where is element \(k\)? Answerable from metadata alone.
  • Packing / semantic layout — the producer’s convention for what each flattened position represents. What is element \(k\)? Not answerable from metadata at all.

A tensor operation can honor the first contract perfectly — move no data, preserve every value — while your interpretation of the second contract is wrong. That’s not a contradiction; the operation did its job. You just decoded the packing with the wrong key.

Here’s the opening problem, made precise. The attention layer packed its output interleaved by head: for head \(h\), role \(\text{qk} \in \{0,1\}\) (0 = query, 1 = key), and feature \(d\), the flat feature index is

\[\text{feature} = h \cdot (2 \cdot d_{\text{head}}) + \text{qk} \cdot d_{\text{head}} + d\]

Stare at that formula and it tells you the axis structure directly: it’s the flattening of a (n_heads, 2, d_head) block, with n_heads the slowest-varying, then role, then feature fastest. Which means the correct way to un-flatten it is view(..., n_heads, 2, d_head) and then pick out role 0 and role 1. Let’s write it, and — crucially — let’s write a second function that we know is correct by brute force, to check the first against.

torch.manual_seed(0)
B, S, n_heads, d_head = 1, 2, 2, 3
D = n_heads * 2 * d_head
# sentinel values so we can read the rearrangement with our eyes, not just trust a bool
x = torch.arange(B*S*D, dtype=torch.float32).reshape(B, S, D)
print("one packed position (batch 0, token 0):", x[0, 0].tolist())
one packed position (batch 0, token 0): [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0]
one packed position (batch 0, token 0): [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0]
def split_qk_reference(x, n_heads, d_head):
    """The oracle. Slow, obviously correct: it transcribes the feature formula directly,
    one head at a time, with zero cleverness. Never ships. Always right."""
    B, S, _ = x.shape
    q = torch.zeros(B, n_heads, S, d_head)
    k = torch.zeros(B, n_heads, S, d_head)
    for h in range(n_heads):
        base = h * 2 * d_head
        q[:, h] = x[:, :, base : base + d_head]              # role 0
        k[:, h] = x[:, :, base + d_head : base + 2*d_head]   # role 1
    return q, k

def split_qk(x, n_heads, d_head):
    """The real implementation: reshape the packed axis into (n_heads, 2, d_head),
    then index out the two roles and move heads ahead of sequence."""
    B, S, _ = x.shape
    x = x.view(B, S, n_heads, 2, d_head)          # decode the packing
    q = x[:, :, :, 0, :].transpose(1, 2)          # role 0 -> (B, n_heads, S, d_head)
    k = x[:, :, :, 1, :].transpose(1, 2)          # role 1
    return q, k

q_ref, k_ref = split_qk_reference(x, n_heads, d_head)
q,     k     = split_qk(x, n_heads, d_head)

print("Q head 0:", q[0, 0, 0].tolist(), " Q head 1:", q[0, 1, 0].tolist())
print("K head 0:", k[0, 0, 0].tolist(), " K head 1:", k[0, 1, 0].tolist())
print("\nmatches the oracle?", torch.equal(q, q_ref) and torch.equal(k, k_ref))
Q head 0: [0.0, 1.0, 2.0]  Q head 1: [6.0, 7.0, 8.0]
K head 0: [3.0, 4.0, 5.0]  K head 1: [9.0, 10.0, 11.0]

matches the oracle? True
Q head 0: [0.0, 1.0, 2.0]  Q head 1: [6.0, 7.0, 8.0]
K head 0: [3.0, 4.0, 5.0]  K head 1: [9.0, 10.0, 11.0]
matches the oracle? True

Because these are integer sentinels and the operation is a pure rearrangement — no arithmetic anywhere — I can use torch.equal for exact shape-and-value equality, which is both correct and stronger than a tolerance-based comparison. (Save torch.testing.assert_close for when actual floating-point math happens and rounding creeps in; a reshape never rounds.) And I can read the result with my eyes: head 0’s query is [0,1,2], its key is [3,4,5] — adjacent in the buffer, exactly as the interleaved packing promised.

Now watch a completely reasonable wrong guess. Suppose you assumed Story B from the opening — roles blocked, not interleaved — and split with chunk:

def split_qk_wrong(x, n_heads, d_head):
    """Assumes packing is [all queries | all keys]. A real convention — just not THIS tensor's."""
    B, S, _ = x.shape
    q_block, k_block = x.chunk(2, dim=-1)                    # split into two halves
    q = q_block.reshape(B, S, n_heads, d_head).transpose(1, 2)
    k = k_block.reshape(B, S, n_heads, d_head).transpose(1, 2)
    return q, k

qw, kw = split_qk_wrong(x, n_heads, d_head)
print("wrong split's shape:", tuple(qw.shape), " same as correct?", qw.shape == q.shape)
print("wrong Q head 0:", qw[0, 0, 0].tolist(), " (correct was [0.0, 1.0, 2.0])")
print("matches the oracle?", torch.equal(qw, q_ref))
wrong split's shape: (1, 2, 2, 3)  same as correct? True
wrong Q head 0: [0.0, 1.0, 2.0]  (correct was [0.0, 1.0, 2.0])
matches the oracle? False
wrong split's shape: (1, 2, 2, 3)  same as correct? True
wrong Q head 0: [0.0, 1.0, 2.0]  (correct was [0.0, 1.0, 2.0])
matches the oracle? False

The wrong split returns the identical shape (1, 2, 2, 3). It raises nothing. Its Q_head0 even happens to coincide, which is worse than if everything differed, because a spot-check on the first head would pass. But it fails the oracle: K_head0 and Q_head1 are scrambled relative to the true packing. If you had no oracle — if you’d trusted the shape and eyeballed one head — this ships, and it degrades your model in a way no stack trace will ever point you toward.

This is the gap between the two contracts, made flesh. The address contract was satisfied by both splits; the packing contract was satisfied by only one; and the machine can only check the first for you. The single most valuable habit here is not a line of code — it’s asking, before you write the split, “how did the producer pack this axis?” Three real conventions show up constantly, and every one is used by code you’ll actually read:

Packing convention How to decode it
[h0_q, h0_k, h1_q, h1_k, …] — interleaved by head view(B, S, n_heads, 2, d_head), index roles [...,0,:] / [...,1,:]
[all_Q ‖ all_K] — blocked by role chunk(2, -1), then view each half to (B, S, n_heads, d_head)
[Q-block ‖ K-block ‖ V-block] — fused QKV, as in nanoGPT view(B, S, 3, n_heads, d_head), index [:, :, 0/1/2]

None is more “correct” than the others; they’re just different promises. Your only job is to find out which promise was made and decode accordingly — and to keep an oracle around so that when you guess, you find out immediately instead of in production.

A second case study: patchify, and a copy you didn’t ask for

The same two-contract thinking cracks open a transform you’ll meet everywhere in vision models: cutting an image into patches and stacking each patch’s pixels into the channel dimension. It’s the operation behind nn.PixelUnshuffle (the downsampling half of many super-resolution networks), and a close cousin of the patch-embedding step that some Vision Transformers use to turn an image into a sequence of tokens. I bring it in here because it’s the cleanest example of a reshape that is forced to copy — the abstract warning from earlier, made concrete and visible.

The goal: take (N, C, H, W) and produce (N, C·p², H/p, W/p), where each p×p spatial block’s pixels get folded into the channels. Watch the shapes at each step:

N, C, H, W, p = 2, 3, 8, 8, 2
img = torch.randn(N, C, H, W)
hp, wp = H // p, W // p

out = (img
       .reshape(N, C, hp, p, wp, p)     # 1. split H into (hp, p) and W into (wp, p)
       .permute(0, 1, 3, 5, 2, 4)       # 2. reorder to (N, C, p, p, hp, wp)
       .reshape(N, C*p*p, hp, wp))      # 3. fold (C, p, p) together into channels

print(tuple(img.shape), "->", tuple(out.shape))
print("matches nn.PixelUnshuffle?", torch.equal(out, torch.nn.PixelUnshuffle(p)(img)))
(2, 3, 8, 8) -> (2, 12, 4, 4)
matches nn.PixelUnshuffle? True
(2, 3, 8, 8) -> (2, 12, 4, 4)
matches nn.PixelUnshuffle? True

Step 1 works because a row index \(r \in [0, H)\) decomposes uniquely as \(r = i \cdot p + \text{ih}\) — patch-row \(i\) (slow) and offset-within-patch \(\text{ih}\) (fast) — which is exactly how reshape splits a contiguous axis. Step 2 slides the within-patch offsets next to the channel axis. Step 3 folds them in. And because a built-in does the same job, PixelUnshuffle is a free oracle — with a bonus check, since its inverse should undo it exactly:

restored = torch.nn.PixelShuffle(p)(out)
print("PixelShuffle(PixelUnshuffle(img)) == img?", torch.equal(restored, img))
PixelShuffle(PixelUnshuffle(img)) == img? True
PixelShuffle(PixelUnshuffle(img)) == img? True

Now the layout point, which is the reason patchify lives in this post rather than a vision one. Look at what the permute in step 2 did to the memory:

after_permute = img.reshape(N, C, hp, p, wp, p).permute(0, 1, 3, 5, 2, 4)
print("after the permute, contiguous?", after_permute.is_contiguous())
print("final reshape shares img's storage?",
      shares_storage(after_permute.reshape(N, C*p*p, hp, wp), img))
after the permute, contiguous? False
final reshape shares img's storage? False
after the permute, contiguous? False
final reshape shares img's storage? False

The permute itself moved nothing — it’s a view, it only changed the stride pattern and thus which axes are logically adjacent. But it left the tensor non-contiguous, and the final reshape then has no choice but to copy: there is no stride assignment that expresses the folded-channel layout over the original buffer, so view would refuse and reshape silently materializes a fresh contiguous tensor. That copy is not a mistake — patchify genuinely rearranges data, and something has to pay for it eventually. But it is a real cost, hiding inside an innocent-looking chain of “just reshaping,” and if it sat in a data pipeline you’d see it as an unexplained dip in throughput with no obvious culprit. Knowing the two contracts is what lets you predict the copy before you profile it.

(The channel ordering that results is c·p² + ih·p + iw — channel-major. That’s a convention; a different permute produces a different, equally valid ordering. If you’re loading someone’s pretrained weights, this is a packing contract you have to match exactly — the same lesson as the Q/K split, wearing a different hat.)

What to carry away

The through-line of everything above is a single reframing: a tensor is not a grid of numbers, it’s a flat buffer plus a promise about how to read it — and really two promises, one the machine checks and one it can’t.

  • The address contract (shape, strides, offset) is pure metadata, and most operations — view, transpose, permute, expand, slicing — just rewrite it, moving no data. From the affine address formula you can derive the rest: why view refuses on non-contiguous tensors, why reshape sometimes copies, why stride-0 expand is free to read but not to write.
  • The packing contract — what each position means — is a convention the producer chose and rarely wrote down. The machine cannot check it. When it’s wrong, your code returns the right shape and the wrong answer, which is the most expensive kind of wrong because nothing announces it.
  • The bridge between them, in practice, is an oracle: a slow, obviously-correct reference you write precisely so that when your fast version guesses the packing wrong, an assertion tells you now instead of a metric telling you next month.

If a single sentence survives: A tensor’s shape describes its logical extent; its strides and offset determine how those coordinates reach storage; and only the producer’s packing contract tells you what those coordinates mean.

Where this goes next

We’ve been treating the buffer as inert storage. The next post asks how the shape rules — broadcasting, reductions, einsum — let you compute across that storage, and how the same “right shape, wrong meaning” trap reappears when two axes happen to share a size. After that: how the GPU actually executes all of this, where the contiguity costs we’ve been hand-waving about get real numbers attached.

References