Most tensor-based code relies on reshape, transpose or permute operations. Some patterns are easy to recognize and understand without having to go into memory layout, strides, view, and contiguity territory. However, when dealing with high-dimensional arrays, these ops are easy to get silently wrong without some kind of mental model and quick intuition about what is going on under the hood.

The goal of this post is to build some layout intuition for these operations: how tensors are structured internally, why they sometimes need to be permuted or why some calls require a copy of the data and some can get away with a view of the original tensor. We will start with the basics and then we will see Grouped-Query Attention as an example.

Some figures are a tad too big, reading this on a phone is not the best experience.

Claude has very politely made all the figures interactive so that you can hover over any cell to follow the element through the panels in the figures.

You can click here if you want to skip to the GQA part.

Tensor internals

A tensor in PyTorch (or NumPy) is not stored as a grid, it’s a single flat strip of numbers plus some metadata. We mostly care about two fields in this metadata: the shape of the tensor and its stride. The shape tells us how many elements we have in each dimension, and the stride tells us how many elements we have to skip in memory to increment that axis index by one.

Let’s get a sample tensor x = torch.arange(40).reshape(4, 2, 5) and say that it represents a (batch, height, width) tensor. The height corresponds to the number of rows, and the width corresponds to the number of columns. The convention that PyTorch follows is row-major order, which means that the last index changes the fastest.

As you can see in the figure below, the reshape puts the first element in the first column, and then places the second one in the second column. This is what it means to “change the fastest”. Only when all the columns in that row have been filled, we increment the height index and continue with the columns in the next row. Once both rows are filled, we continue with the next batch element.

This links to the concept of strides we have briefly mentioned earlier. For a contiguous tensor (I’m getting a bit ahead of myself here), each stride is the product of the sizes of the dimensions to its right. As an example, for x with shape (4, 2, 5), the strides are (10, 5, 1). As mentioned above, the strides tell us how many elements we have to “skip” in order to increase that axis index by one. In the case of our batch dimension, we need to skip 10 elements. The caption of the figure shows this computation for every element in the tensor. Mathematically:

The two parts of the formula are not equally general. The address sum holds for any strided tensor, the product only gives the strides of a contiguous one.

This formula is the mental model we will try to internalize, visualizing how reshape and permute change either how we group the elements or the strides of each axis.

Reshaping or Viewing a tensor

When we reshape/view a tensor we are not doing anything other than taking the internal layout of the tensor and defining how many dimensions we want and how many elements we want in each dimension. We are regrouping the strip.

view and reshape are similar. The difference is that if the tensor has incompatible strides, reshape will create a contiguous copy before regrouping, while view will raise an error.

Permuting or Transposing a tensor

permute changes the order of the axes. Let’s look at what y = x.permute(1, 0, 2) is doing:

• Axis 1 of x becomes axis 0 of y.
• Axis 0 of x becomes axis 1 of y.
• Leaves axis 2 in the same place.

transpose is just a shortcut to swap two axes. Considering the previous operation: x.permute(1, 0, 2) == x.transpose(0, 1) == x.transpose(1, 0).

Two dimensions: easy peasy, we saw this in school

Permuting two axes is what most people understand by the “transpose” of a matrix. If we imagine a 2D matrix, transposing it means to mirror it across the diagonal, swapping the (i, j) element with the (j, i) element.

In the example below, we see that we are swapping the dims, but also the strides. The addresses in memory don’t change, only their arrangement in the tensor and how we read them. A very important point: you might have noticed that by swapping the strides together with the dimensions, we have broken the rule that defines a contiguous tensor, the strides are no longer the product of the dimensions to the right, so we no longer have a contiguous tensor.

Three dimensions: a bit less easy

We could visualize this as a 3D cube in space, but I want to use this “easier” case as a way to start visualizing N dimensions in a plane. If we have three dimensions, we can keep the last two as rows and columns, just like before, and place the “batch elements” (each index of the first axis) next to each other.

The following figure shows all the possible permutations in a three-axis tensor. Pay special attention to how the elements of the output tensor change when we swap the axes. The shape is “easy” to predict, but the order of the elements is a bit less intuitive.

(0, 1, 2) is trivial, (0, 2, 1) is just a 2D transpose in each batch element.

(1, 0, 2) and (2, 1, 0) also swap two axes, we can still use the “mirroring” visualization, but this time around elements that are not visually together in our mental layout. Instead of trying to mirror the tensor, let’s try to use these as a stepping stone to more complex permutes.

For the (1, 0, 2) case, let’s read the permute as a mapping: axis 0 of y is axis 1 of x, axis 1 of y is axis 0 of x, and axis 2 stays where it was. So all items sharing a row index in x will now share a batch index in y: all the blue elements (row 0 in x) fall in the first batch element of y, and all the orange ones (row 1) fall in the second. Inside each new batch element, the rows are the old batch elements, so we “pick” the row with the same index from every batch element of x and stack them. The shape follows the same mapping, (4, 2, 5) becomes (2, 4, 5), which is a good sanity check. (2, 1, 0) is also a two-axis swap that involves the batch axis, you can try to predict the output layout.

Both (1, 2, 0) and (2, 0, 1) permute all three axes at the same time. Let’s leave the stride arithmetic aside and use the same mapping to get the layout “visually”.

In the (1, 2, 0) case: axis 0 of y is axis 1 of x, so each batch element of y groups the items that shared a row index in x (the blue ones in one, the orange ones in the other); axis 1 of y is axis 2 of x, so each row of y holds the items that shared a column index in x; and axis 2 of y is axis 0 of x, so moving along a row of y walks through the old batch elements. Again, the shape (2, 5, 4) confirms it. The reasoning for (2, 0, 1) is left as an exercise for the reader.

More dimensions: spooky

For a larger number of dimensions, I find it easier to stop thinking in rows and columns and give each axis a spatial meaning: place it where what it measures would “physically” be. Doing this makes the “row-major” intuition a bit brittle, so we have to be careful to walk the axes from right to left when flattening the memory layout to, for example, reshape the tensor.

The figures below show an image being cut into patches for a ViT (without using a Conv2D op). If we take a C, H, W = (3, 6, 6) image and cut it into 3x3 patches, we should end up with 4 patches of \(3 \times 3 \times 3 = 27\) elements.

img.shape
# (3, 6, 6): channel, row, col
cut = img.reshape(3, 2, 3, 2, 3)
# channel, grid row, patch row, grid col, patch col
moved = cut.permute(1, 3, 0, 2, 4)
# grid row, grid col, channel, patch row, patch col
patches = moved.reshape(4, 27)
# patch, value

We start with the image as three blocks, one per channel:

We reshape the image from (channels, rows, cols) into (channels, grid_rows, patch_rows, grid_cols, patch_cols). The grid axes say which patch a pixel belongs to, and the patch axes say where it sits inside that patch. The figure below shows what I mean: inside each channel, the grid axes lay the patches out as a \(2 \times 2\) grid over the image, and the patch axes lay the pixels out inside each patch, exactly where they were. Take into account that you are now working with a 5-D tensor (!).

As we saw before, the memory remains untouched: the reshape only changed the shape and the strides. Now cut.permute(1, 3, 0, 2, 4) reorders the axes into (grid row, grid col, channel, patch row, patch col). We can place each axis in the drawing however we want: the tensor has no notion of “across” or “down”. Nonetheless, we cannot forget their order (the numbers on the arrows), because that’s what decides how a later reshape flattens them in row-major order. Going over the axes from right to left:

• Axis 4, patch col: unchanged, three pixels side by side.
• Axis 3, patch row: the three rows of a patch are now placed next to each other, so one channel of a patch becomes a single row of nine pixels.
• Axis 2, channel: the three channels of the same patch are stacked on top of each other. Each patch is now a \(3 \times 9\) block holding exactly the 27 values that will end up in its row of patches.
• Axes 0 and 1, grid row and grid col: these only say which patch we are looking at. I placed grid row going down and grid col going across, so every patch stays where it was in the image. If we forget about the other three axes, we are looking at a \(2 \times 2\) tensor where each element is a whole patch.

For completeness, let’s apply the final reshape we would perform before the linear projection. Since we have scrambled the axes and made the tensor non-contiguous, PyTorch performs a copy of the data and reorders its internal representation before reshaping (see next section for more info on this operation). We reshape to (grid_rows * grid_cols, channels * patch_rows * patch_cols) == (4, 27). We have to be careful here with the mental image, we have laid out the axes in a way that was visually convenient for us in “conceptual” terms for the patches, but when flattening the data to reorganize it the reshape follows row-major ordering. Make sure that the axes that are being merged are “neighbors” and go from slowest to fastest.

Making a tensor contiguous

In all the permute figures we have kept the memory strip at the bottom the same. As we have already mentioned a couple of times, permute doesn’t change the numbers at all, it just shuffles the shapes and the corresponding strides. The formula for the memory address still works, we just change how we walk over the strip, essentially jumping over different “sections” of it. Going back to our simpler example, let’s take a look at what this walking looks like:

Once again, this is what PyTorch calls a non-contiguous tensor: the memory order and the layout have come apart. reshape will work fine with a non-contiguous tensor, but view never copies, so it may raise an error: it only works if the new shape can be read from the same strip just by picking new strides. To be more precise, each new dimension of the view must either:

• Be a subspace of a single original dimension (e.g. unflattening a dimension), or
• Span a set of original dimensions \(d, d+1, ..., d+k\) that, for all \(i = d, ..., d+k-1,\) satisfy \(\text{stride}[i] = \text{stride}[i+1] \times \text{shape}[i+1]\).

For example, y = x.permute(1, 0, 2) has shape (2, 4, 5) and strides (5, 10, 1). It’s non-contiguous, but y.view(2, 2, 2, 5) works, since it only splits the axis of size 4 into two axes with strides (20, 10). y.view(8, 5) raises an error: it would merge the first two axes, but \(5 \neq 10 \times 4\), so no single stride can walk over both of them in order.

Under the hood, reshape performs (if needed) a torch.Tensor.contiguous() which basically returns itself if already contiguous, or allocates a copy and writes the data in contiguous order. This is shown in the figure below.

This is the only operation in this post that moves the underlying data (copies, actually).

Grouped-Query Attention

As promised, let’s go over a more practical example now. Regular multi-head attention has the same number of heads for the queries, keys and values (n_heads). Grouped-Query Attention keeps n_heads for the queries but gives a smaller number of heads to the keys and values (n_kv_heads), with each key/value head being shared by q_per_kv_head = n_heads / n_kv_heads query heads. The main advantage of this is shrinking the KV cache at inference time.

Setup

For the example I’ll use small numbers to excuse the reader from needing to buy an ultrawide monitor just to read the post. The dimensions and colors for each axis are as follows:

The six query heads get six hues: heads 0 to 2 are the cool ones and share the blue key head, heads 3 to 5 are the warm ones and share the orange one.

I will leave the batch dimension out of the figures but keep it in the shapes and the code.

For reference, the code (ignoring causal masking) should be something along these lines:

# x: (B, T, d_model)
q_per_kv_head = n_heads // n_kv_heads  # 3

q = wq(x)  # (B, 4, 30)
k = wk(x)  # (B, 4, 10)
v = wv(x)  # (B, 4, 10)

q = q.reshape(B, T, n_kv_heads, q_per_kv_head, d_head)
q = q.permute(0, 2, 3, 1, 4)  # (B, 2, 3, 4, 5)
k = k.reshape(B, T, n_kv_heads, 1, d_head)
k = k.permute(0, 2, 3, 1, 4)  # (B, 2, 1, 4, 5)
v = v.reshape(B, T, n_kv_heads, 1, d_head)
v = v.permute(0, 2, 3, 1, 4)  # (B, 2, 1, 4, 5)

att = q @ k.transpose(-2, -1) / d_head ** 0.5  # (B, 2, 3, 4, 4)
att = att.softmax(dim=-1)
out = att @ v  # (B, 2, 3, 4, 5)

out = out.permute(0, 3, 1, 2, 4)  # (B, 4, 2, 3, 5)
out = out.reshape(B, T, n_heads * d_head)  # (B, 4, 30)

Projecting the inputs

The three projections give us one row per token: 30 values for the queries (6 heads of 5 features) and 10 for the keys and the values (2 heads of 5 features). Values look exactly like the keys, so I’m only drawing the keys.

Splitting into heads

The reshape splits those 30 values into three axes: which key/value head the query head belongs to, which of the query heads of that group it is, and the feature. Going back to giving each axis a spatial meaning, here we keep the tokens going down like a sequence and lay the three new axes across, from the biggest steps to the smallest: the two groups side by side (one per key/value head), the three query heads inside each group, and the five features inside each head. The keys get the same treatment, with a 1 where the queries have q_per_kv_head.

As we already know, it’s a reshape, so nothing moved, the rows of the previous figure just got regrouped.

The order of the two new axes is a decision: (n_kv_heads, q_per_kv_head) puts query heads 0, 1 and 2 with key/value head 0 and heads 3, 4 and 5 with key/value head 1. Writing (q_per_kv_head, n_kv_heads) would pair heads 0, 2 and 4 with key/value head 0 (both are valid conventions for the same architecture: interleaved vs tiled, be careful if loading weights).

Preparing for the matmuls

permute(0, 2, 3, 1, 4) keeps the batch at the front and takes the token axis all the way to the back, next to the features. As we will see in the next subsection, this allows for batched matrix multiplication of queries and keys.

The drawing doesn’t need to change at all, this time we change the numbers in the arrows (the axis index on each arrow). As we have seen in the first half of the post, this matters for reshaping, but we can lay them down visually however we want.

Now, the last two dimensions are (token, feature), so every head is a small matrix with one row per token and one column per feature, and the two axes in front only say which matrix we are looking at: which group, and which head inside the group. The keys are two matrices, one per group.

Computing the scores and broadcasting

@ only looks at the last two axes and treats everything in front as a batch, so it works matrix by matrix: the \(4 \times 5\) matrix of each query head gets multiplied with the matrix of its key head. For the shapes to fit, the keys need their last two axes swapped. transpose(-2, -1) is the mirror we have seen previously applied inside every matrix, which leaves them with one row per feature and one column per token.

Axis 2 has size 3 in the queries and size 1 in the keys, so the single key matrix of each group is read once for each of the three query heads of that group (broadcasting). In the figure those extra readings are the dotted copies.

After the softmax, att @ v broadcasts in the same way and returns a tensor shaped like the queries, (B, 2, 3, 4, 5), i.e. (B, n_kv_heads, q_per_kv_head, T, d_head).

Merging the heads back

The output has to go back to one row of 30 values per token. permute(0, 3, 1, 2, 4) undoes the first permute and brings the tokens back to axis 1, leaving the three axes we want to merge at the end: key/value head, query head inside the group, feature (in order).

Once again the picture is the same and only the numbers on the arrows moved to the order we had right after splitting the heads. Now the last reshape can merge the three axes at the end (closing the horizontal gaps in the figure).

Since the permuted tensor is not contiguous, that reshape has to make a contiguous copy. The figure below illustrates the contiguous() we have already seen for our particular case: every set of five numbers (d_head) moves from its place in the old strip, where the memory goes head after head, to its place in the new one, where it goes token after token.

And we have our (B, T, n_heads * d_head) back!

Just for reference, compare this to doing the reshape without the permute, which gives a valid shape but messes up the structure inside the tensor:

In this case for example, the first row holds head 0 for all four tokens followed by head 1 for tokens 0 and 1.

Summary

• A tensor is a 1D array plus a header. The header holds the shape and the strides, and the address of an element is \(\sum{\text{index}_i \cdot \text{stride}_i}\).
• reshape and view regroup elements. Any shape with the same number of elements works (read in row-major order), but the result only keeps its meaning if we merge or split neighboring axes in order. They don’t change the memory order unless reshape performs a contiguous under the hood.
• permute and transpose reorder axes. They shuffle shape and strides together but don’t move anything.
• contiguous moves numbers. We need it between a permute and a view, reshape does it when needed.
• With many axes forget about rows and columns, name every axis by what it measures, give it a spatial meaning, and check their order before reshaping.

Thanks for reading!

Related links

• PyTorch internals (the part about strides)
• torch.Tensor.view
• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints