Fold the heads into the batch

Some attention operations expect one batch axis. If activations have shape (batch, heads, length, dim), you can combine batch and head indices:

(batch, heads, length, dim) -> (batch * heads, length, dim)

The new batch index b * heads + h identifies the original pair (b, h). Every head keeps its sequence and feature values in the same order.

When this can share storage

For a contiguous input, merging these leading axes requires no data copy:

import torch

x = torch.arange(120).reshape(2, 3, 4, 5)
y = x.reshape(6, 4, 5)
print(x.stride())  # (60, 20, 5, 1)
print(y.stride())  # (20, 5, 1)
print(torch.equal(y[4], x[1, 1]))  # True
print(y.untyped_storage().data_ptr() == x.untyped_storage().data_ptr())  # True

Each step along the merged axis advances 20 elements. The original batch stride, 60, equals three head strides. This lets a single stride describe the combined axis. A view still needs tensor metadata, but it reuses the existing data buffer.

Adjacent axes are not enough

Strides matter, not just where the dimensions appear in the shape. Starting from the example above:

z = x.permute(0, 2, 1, 3)  # shape (2, 4, 3, 5)
print(z.stride())  # (60, 5, 20, 1)
folded = z.reshape(8, 3, 5)
print(folded.untyped_storage().data_ptr() == x.untyped_storage().data_ptr())  # False

Here the leading axes cannot be merged into a view: moving between batches requires a different spacing from moving between sequence positions. reshape copies to preserve the requested element order. This example uses non-singleton dimensions; layouts with size-one axes can behave differently.

view raises when the requested shape cannot share storage. It does not require every input to be contiguous: some noncontiguous layouts still support a particular view. reshape accepts either case and copies when necessary.

Give each tensor axis a meaning

A tensor's shape tells you how many entries each axis has. It does not tell you what those entries represent. For example, (batch, features) and (features, batch) use different axis orders. Name the axes before choosing an operation.

A reduction combines entries along an axis, for example by summing or averaging. Averaging the features of a (B, D) tensor gives one number per example, with shape (B,). Keeping the reduced axis gives shape (B, 1) instead.

For [[2, 4, 6], [1, 3, 5]], the row means with the axis kept are [[4], [3]]. Subtracting them from the input gives [[-2, 0, 2], [-2, 0, 2]]: each row has been centered using its own mean.

Reshape and transpose do different jobs

A reshape regroups elements using the operation's traversal order. A transpose changes the axis order. For the usual row-major reshape, [[1, 2, 3], [4, 5, 6]] reshaped to (3, 2) becomes [[1, 2], [3, 4], [5, 6]]. Transposing the original instead gives [[1, 4], [2, 5], [3, 6]].

Both results have shape (3, 2), but their entries mean different things. Check values as well as shapes. Unequal axis sizes make many axis mistakes easier to spot than square examples do.

Concatenate along an axis or stack on a new one

Concatenation extends an existing axis. Stacking creates a new axis. For two arrays of shape (2, 3), concatenating on axis zero gives (4, 3). Stacking on a new leading axis gives (2, 2, 3).