Reshape vs transpose

A tensor’s numbers sit in one long line. Row-major order (PyTorch’s default) stores row 1, then row 2: the last axis changes fastest.

reshape (or view) keeps the flat order and cuts it into new rows: 1 2 | 3 4 | 5 6. The numbers don’t move in memory.

The transpose is also $3 \times 2$, but its rows are $[1, 4], [2, 5], [3, 6]$. Same shape as the reshape, different matrix. PyTorch transposes by swapping the strides (how far to jump in memory for each axis), so nothing is copied, and the result is no longer laid out in row-major order: it’s non-contiguous.

Two tokens, $d = 4$ features, $h = 2$ heads of $d_h = 2$. view(2, 2, 2) cuts each token’s 4 features into 2 heads: head 0 gets features 1–2, head 1 gets features 3–4. Shape $(T, h, d_h)$. Correct so far.

transpose(0, 1) swaps the token and head axes: $(h, T, d_h)$. Each head now holds one row per token, made of that token’s own slice of features. Attention can run on each head separately.

view(2, 2, 2) read as $(h, T, d_h)$ cuts the flat list a₁ a₂ a₃ a₄ b₁ b₂ b₃ b₄ in the wrong places. Head 0’s “two tokens” are really two halves of token a; token b vanished from it. The shape is right, so no error, but attention now mixes features as if they were tokens.

q = x.view(B, T, h, d_h).transpose(1, 2)      # (B, h, T, d_h)
# ... attention per head ...
out = out.transpose(1, 2).reshape(B, T, h * d_h)   # (B, T, d), heads merged

After transpose, the tensor is non-contiguous, so .view would fail. .reshape copies when it has to (or write .contiguous().view(...)).

With $h = 3$ heads, write the shape of x.view(2, 6, 3, 4).transpose(1, 2), and what each axis means.

The view gives $(2, 6, 3, 4)$ = (batch, token, head, feature); swapping axes 1 and 2 puts heads before tokens.

• Memory is flat and row-major: the last axis changes fastest.
• reshape/view re-cuts the flat order; transpose swaps axes (strides) and makes the tensor non-contiguous.
• Split heads: view(B, T, h, d_h).transpose(1, 2). Viewing straight to $(B, h, T, d_h)$ scrambles tokens without an error.
• Merge back: transpose(1, 2).reshape(B, T, d).