NumPy and Tensor Shapes

Two short text snippets contain three words each: "the red fox" and "a blue bird." Suppose we represent each word with four numbers. We have 24 numbers in total, but the total alone won't tell us which numbers belong to which word.

The previous Python for AI Engineering lesson scored JSONL rows one at a time: parse a line, check its fields, compare a predicted status with the expected one. Here, the same kind of data arrives as one numeric array holding many rows, and one operation runs across them together.

Two log snippets, one array

An array arranges values along numbered axes, such as rows and columns. NumPy calls its array object an ndarray. In machine-learning code, you'll also see the word tensor for an array with zero, one, or several axes.

A token is a piece of text a model processes; here we choose one word per token. A feature vector is the list of numbers representing one token. The numbers below are deliberately artificial, not learned word meanings: counting from 0 makes it easy to trace every value when we rearrange the array.

Give each token four features and fill them in order so every coordinate has a real value. Before reading the shape, count the snippets, token rows, and feature columns:

x =
[
  [[ 0,  1,  2,  3],
   [ 4,  5,  6,  7],
   [ 8,  9, 10, 11]],

  [[12, 13, 14, 15],
   [16, 17, 18, 19],
   [20, 21, 22, 23]]
]

There are two snippets, three token rows per snippet, and four values per row. So the shape is (B, T, D) = (2, 3, 4):

If you can say that sentence, you can usually debug the next operation. The first axis picks a snippet, the second picks a word, and the third holds that word's four numbers. x[0, 1] selects the row labeled red and returns its four-feature vector.

View larger

Rank, shape, and dtype

Finding one value takes three indices, and the shape says how many choices exist at each position. The number of axes is the array's rank, exposed as x.ndim, so this array has rank 3. Axis numbers start at zero: batch is axis 0, tokens are axis 1, and features are axis 2. Negative axis numbers count backward, making -1 another name for the feature axis.

The comma in (4,) is Python's notation for a one-item tuple. An array's shape is always a tuple, even when it contains just one axis length.

Shape is only half of the contract. Every NumPy array also has one data type, or dtype, that says how each element is stored and interpreted. The examples use float32, a 32-bit floating-point type common in model inputs. np.arange(24) without a dtype uses NumPy's default integer; floating-point arguments can select a floating-point dtype instead. Writing dtype=np.float32 makes later arithmetic match the type you'll see in model code. Shape tells you where values live; dtype tells you what kind of values they are.

At a model boundary, check shape and dtype separately. An array with shape (2, 3, 4) and dtype int64 passes the shape check but fails a contract that requires float32.

Indexing: which axes stay and which disappear

Each coordinate now points at an actual token vector.

Before using the table, predict x[:, 1].shape: the batch axis stays, the token axis is selected, and the four features stay, so (2, 4) is the answer.

Read an index from left to right. A single integer such as 0 or 1 chooses one item and removes that axis. A colon keeps the whole axis. A slice such as 0:1 also keeps the axis, but selects only its first item: x[0:1] has shape (1, 3, 4), unlike x[0] with shape (3, 4). Use the slice when later code still expects a batch axis.

Inspect the batch

Run the script to check both values and shapes. np.arange(24) creates the integers from 0 through 23; reshape(B, T, D) groups them into the nested layout above. With uv from the Python lesson, save a cell as a .py file and run uv run --with numpy filename.py.

import numpy as np

B, T, D = 2, 3, 4
x = np.arange(B * T * D, dtype=np.float32).reshape(B, T, D)

print("x shape", x.shape)
print("x dtype", x.dtype)
print("x[0] shape", x[0].shape)
print("x[0,1]", x[0, 1])
print("x[:,1]")
print(x[:, 1])

x shape (2, 3, 4)
x dtype float32
x[0] shape (3, 4)
x[0,1] [4. 5. 6. 7.]
x[:,1]
[[ 4.  5.  6.  7.]
 [16. 17. 18. 19.]]

Selecting the token red gave us four features. Adding the same four offsets to every token should preserve those feature positions, without writing a loop over snippets and tokens.

Broadcasting: reuse one feature bias

A bias is an additive offset. Here use four offsets, one per feature, and add them to every token vector. Should that addition keep all three axes or lose one?

bias = [100, 200, 300, 400]

This bias has shape (4,), so it matches the last axis of x.

For the red token, the four additions are:

x[0, 1]     = [  4,   5,   6,   7]
bias        = [100, 200, 300, 400]
result      = [104, 205, 306, 407]

NumPy applies that addition to each token in each snippet without a Python loop. You can picture the bias vector repeating across the batch and token axes, but that repetition is conceptual: broadcasting avoids materializing a full repeated bias array here. The addition still allocates the full (2, 3, 4) result, so broadcasting doesn't make a large output free.^{[1]Reference 1Broadcasting.https://numpy.org/doc/stable/user/basics.broadcasting.html}

The batch and token axes stay. The feature bias repeats across both of them, so the output keeps shape (2, 3, 4).

The matching rule:

1. Line shapes up from the right.
2. Two axes are compatible if they're equal or one of them is 1.
3. Missing leading axes in the smaller array count as size 1 (NumPy prepends them mentally).
4. If neither rule holds, NumPy raises an error.

For x.shape == (2, 3, 4):

Each column below represents a fixed axis position. Cyan cells are size-1 axes, including the missing leading axes. The failure puts 3 on D, not on T.

View larger

Add the feature bias

Predict the selected row after the addition: [104, 205, 306, 407]. Then run the cell and check that its shape is still (2, 3, 4).

import numpy as np

B, T, D = 2, 3, 4
x = np.arange(B * T * D, dtype=np.float32).reshape(B, T, D)
bias = np.array([100, 200, 300, 400], dtype=np.float32)
shifted = x + bias

print("shifted shape", shifted.shape)
print("shifted[0,1]", shifted[0, 1])

shifted shape (2, 3, 4)
shifted[0,1] [104. 205. 306. 407.]

If you accidentally try np.array([10, 20, 30]), NumPy raises:

ValueError: operands could not be broadcast together with shapes (2,3,4) (3,)

That message is useful. It shows the feature axis wasn't where you thought it was.

You can test that failure deliberately instead of waiting for it to surprise you:

import numpy as np

x = np.arange(2 * 3 * 4, dtype=np.float32).reshape(2, 3, 4)
token_bias = np.array([10, 20, 30], dtype=np.float32)

try:
    x + token_bias
except ValueError as exc:
    print(str(exc))
else:
    raise RuntimeError("Expected NumPy to reject shape (3,) against (2, 3, 4)")

operands could not be broadcast together with shapes (2,3,4) (3,)

When broadcasting succeeds with the wrong meaning

A 1-D array (3,) isn't the same as a column-shaped array (3, 1). Inside an index, None inserts an axis of size 1 without changing the values. Here's how to make either orientation explicit:

import numpy as np

v = np.array([1, 2, 3])
print("v shape", v.shape)
print("v[:, None] shape", v[:, None].shape)
print("v[None, :] shape", v[None, :].shape)

v shape (3,)
v[:, None] shape (3, 1)
v[None, :] shape (1, 3)

A plain (3,) has no row or column orientation. Broadcasting aligns its only axis with the trailing axis, which may not be the axis you intended.

Now temporarily take just the first three features of snippet 0. Its shape is (3, 3), so token count and feature count happen to match. NumPy will accept a (3,) token bias, but apply it to features instead. Compare the first row under both operations:

import numpy as np

x = np.arange(24, dtype=np.float32).reshape(2, 3, 4)
square = x[0, :, :3]
token_bias = np.array([10, 20, 30], dtype=np.float32)
wrong = square + token_bias
right = square + token_bias[:, None]

print("both shapes", wrong.shape, right.shape)
print("wrong first token", wrong[0])
print("right first token", right[0])
assert np.array_equal(right[0], [10, 11, 12])
assert not np.array_equal(wrong, right)

both shapes (3, 3) (3, 3)
wrong first token [10. 21. 32.]
right first token [10. 11. 12.]

The intended offset for token 0 is 10, applied to all three features. A successful operation and a correct output shape don't prove that the axes mean what you intended. Check a known row's values too.

Reductions: collapsing the right axis

A reduction answers the question, "which axis am I summarizing away?"

Suppose you want one vector per snippet, not one vector per token. Then average across the token axis. Predict the result before writing the call: keeping batch and feature axes should leave shape (2, 4).

x.mean(axis=1)

Keep using the same x whose values run from 0 through 23:

For snippet 0, average each feature column across its three token rows:

token 0 = [ 0,  1,  2,  3]
token 1 = [ 4,  5,  6,  7]
token 2 = [ 8,  9, 10, 11]
mean    = [ 4,  5,  6,  7]

So x.mean(axis=1) keeps the batch and feature axes, giving shape (B, D) = (2, 4). The token axis disappears; one vector remains for each snippet.

Now compare that with:

x.mean(axis=-1)

That reduces the feature axis instead. It keeps batch and token axes, so the shape becomes (B, T) = (2, 3).

Compare two reductions

The next cell makes the contrast visible. Predict (2, 4) for token pooling and (2, 3) for feature means, then check the first row values.

import numpy as np

B, T, D = 2, 3, 4
x = np.arange(B * T * D, dtype=np.float32).reshape(B, T, D)
pooled_tokens = x.mean(axis=1)
feature_means = x.mean(axis=-1)

print("pooled_tokens shape", pooled_tokens.shape)
print("pooled_tokens dtype", pooled_tokens.dtype)
print("pooled_tokens[0]", pooled_tokens[0])
print("feature_means shape", feature_means.shape)
print("feature_means[0]", feature_means[0])

pooled_tokens shape (2, 4)
pooled_tokens dtype float32
pooled_tokens[0] [4. 5. 6. 7.]
feature_means shape (2, 3)
feature_means[0] [1.5 5.5 9.5]

Before writing a reduction, say its contract out loud: "I want one vector per snippet, so I'm reducing the token axis." That sentence tells you which axis disappears and which two must remain.

The float32 input produces a float32 mean. By default, integer inputs use float64 accumulation and output; floating-point inputs return their input dtype. float16 is a special case: NumPy accumulates in float32 for extra precision before returning float16.^{[2]Reference 2numpy.meanhttps://numpy.org/doc/stable/reference/generated/numpy.mean.html}

A large float32 reduction can lose precision. Pass dtype=np.float64 when you need a higher-precision accumulator and can accept a float64 result. Check both .shape and .dtype when a reduction feeds later model code.

Even three values can expose the problem. Suppose one feature across three tokens is [16777216, 1, -16777216]. In exact arithmetic, the sum is 1, so the mean is 1/3. At that magnitude, float32 can't retain the extra 1 when it adds the first two values. Compare the same stored inputs with two accumulator dtypes:

import numpy as np

feature = np.array([16777216, 1, -16777216], dtype=np.float32)
default_mean = feature.mean()
wide_mean = feature.mean(dtype=np.float64)

print("float32 mean", float(default_mean))
print("float64 mean", float(wide_mean))
assert default_mean == 0
assert np.isclose(wide_mean, 1 / 3)

float32 mean 0.0
float64 mean 0.3333333333333333

All three inputs are exactly representable here. The lost unit comes from accumulation, so a wider accumulator recovers it. It won't recover information already rounded away when creating the input array. The values from 0 through 23 in our running batch don't trigger this failure.^{[2]Reference 2numpy.meanhttps://numpy.org/doc/stable/reference/generated/numpy.mean.html}

Subtract a mean without losing its position

By default, a reduction removes the axis entirely. That can break later broadcasting because the shape has changed. Predict which result can subtract back into (2, 3, 4): (2, 4) or (2, 1, 4)?

import numpy as np

B, T, D = 2, 3, 4
x = np.arange(B * T * D, dtype=np.float32).reshape(B, T, D)

# Without keepdims: shape becomes (2, 4), which no longer broadcasts with (2, 3, 4)
mean_flat = x.mean(axis=1)
print("mean_flat shape", mean_flat.shape)

# With keepdims: shape stays (2, 1, 4), which still broadcasts with (2, 3, 4)
mean_kept = x.mean(axis=1, keepdims=True)
print("mean_kept shape", mean_kept.shape)
centered = x - mean_kept
print("centered snippet 0")
print(centered[0])
assert np.allclose(centered.mean(axis=1), 0)

mean_flat shape (2, 4)
mean_kept shape (2, 1, 4)
centered snippet 0
[[-4. -4. -4. -4.]
 [ 0.  0.  0.  0.]
 [ 4.  4.  4.  4.]]

keepdims=True keeps the collapsed axis as a placeholder of size 1. That preserves the original axis positions so operations against the original tensor can still broadcast correctly. Use it when a reduction result must be added, subtracted, or divided back into its source-shaped tensor.

The first feature now contains [-4, 0, 4] across tokens, whose mean is zero. np.allclose checks that property with a tolerance for floating-point rounding. If batch size happened to equal token count, the version without keepdims could run and subtract the wrong snippet's mean, just like the square-array bias bug.

Side by side, the missing token slot is the whole bug. Without keepdims, (2, 4) right-aligns so the leftover 2 lands on T. With keepdims, a cyan size-1 slot stays where T was, and the mean can subtract back into x.

View larger

Reshape vs. transpose: regrouping or reordering

x still means batch, then tokens, then features. A later line asks for shape (2, 4, 3). Is that the same 24 numbers regrouped, or the same axes moved?

Use the same batch so you can see how identical values acquire different positions:

import numpy as np

x = np.arange(24, dtype=np.float32).reshape(2, 3, 4)
reshaped = x.reshape(2, 4, 3)
transposed = np.transpose(x, (0, 2, 1))

print("x[0]")
print(x[0])
print("reshaped[0]")
print(reshaped[0])
print("transposed[0]")
print(transposed[0])

x[0]
[[ 0.  1.  2.  3.]
 [ 4.  5.  6.  7.]
 [ 8.  9. 10. 11.]]
reshaped[0]
[[ 0.  1.  2.]
 [ 3.  4.  5.]
 [ 6.  7.  8.]
 [ 9. 10. 11.]]
transposed[0]
[[ 0.  4.  8.]
 [ 1.  5.  9.]
 [ 2.  6. 10.]
 [ 3.  7. 11.]]

Both results have shape (2, 4, 3), but they answer different questions. Predict where the highlighted 3 should land while it stays attached to token the, feature d3. reshape parks it next to 4, the first feature of token red; transpose keeps it at feature d3 of the in a feature-by-token grid.

The figure also shows where the values live in memory. Each float32 occupies four bytes. For x[0], moving one feature takes four bytes, while moving one token takes 4 × 4 = 16 bytes. These byte steps are its strides, (16, 4). Shape counts positions along each axis; strides describe how to reach those positions in the buffer.^{[3]Reference 3numpy.ndarray.strideshttps://numpy.org/doc/stable/reference/generated/numpy.ndarray.strides.html}

Our original array is C-contiguous: the last axis advances through neighboring stored values, then the next axis advances to the next row. Reshaping this snippet to (4, 3) changes its row length to three values, so its strides become (12, 4). Transposing it swaps the existing strides to (4, 16). The values stay in the same buffer, but stepping across a transposed row now skips 16 bytes at a time.

View larger

If your real goal is to turn (B, T, D) into (B, D, T), that's an axis-order problem, so transpose is the right tool.

Reshaping must preserve element count: 2 * 3 * 4 and 2 * 4 * 3 both equal 24. One dimension can be -1 to let NumPy infer its size. For example, x.reshape(-1, 4) gives (6, 4), combining batch and token positions while keeping each four-feature row intact.^{[4]Reference 4numpy.reshapehttps://numpy.org/doc/stable/reference/generated/numpy.reshape.html}

A 1-D vector has a quieter trap. Predict both printed shapes before running the cell: transposing (3,) still gives (3,).

import numpy as np

v = np.array([1, 2, 3])
print(v.shape)
print(v.T.shape)

(3,)
(3,)

A 1-D vector has one axis, so transposing it does nothing. If you need a row or column vector, make that axis explicit with v[None, :] or v[:, None].

Views and copies: who shares memory

An ndarray combines a buffer of values with metadata such as shape, dtype, and strides. A view reuses a buffer with its own metadata; a copy owns separate values.

Basic slicing, such as x[:, 1, :], returns a view. Selecting with an integer array or boolean mask is advanced indexing, which returns a copy.^{[5]Reference 5Copies and views.https://numpy.org/doc/stable/user/basics.copies.html}

np.transpose returns a view whenever possible; for the ndarray below, it shares the buffer and changes the strides. reshape returns a view when the new layout can be expressed with strides, but it may copy a non-contiguous array.^{[5]Reference 5Copies and views.https://numpy.org/doc/stable/user/basics.copies.html}

Predict the mutation before running the cell: editing the basic slice should change x, while editing the advanced-index result should leave x alone.

import numpy as np

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

token = x[:, 1, :]   # basic slice -> view shares memory with x
token[0, 0] = 99
print("after editing a slice, x[0, 1, 0] =", x[0, 1, 0])

picked = x[[0, 1]]   # advanced indexing -> independent copy
picked[0, 0, 0] = -1
print("after editing an advanced-index copy, x[0, 0, 0] =", x[0, 0, 0])

after editing a slice, x[0, 1, 0] = 99.0
after editing an advanced-index copy, x[0, 0, 0] = 0.0

If you need to mutate a slice without touching the source, take an explicit copy first with x[:, 1, :].copy(). When unsure, np.shares_memory(a, b) answers the question directly.

Retrieving an advanced-index result and assigning through advanced indices are different operations. picked = x[[0, 1]] creates a copy, but x[[0, 1]] = 0 writes directly into x. The copy rule doesn't make indexed assignment harmless.^{[5]Reference 5Copies and views.https://numpy.org/doc/stable/user/basics.copies.html}

A transpose can make the next reshape copy

A transpose changes how indices traverse the buffer. Its result is often no longer C-contiguous, so a later reshape may need to copy values into a new buffer. The code checks strides and sharing before and after that reshape:

import numpy as np

x = np.arange(24, dtype=np.float32).reshape(2, 3, 4)
t = x.transpose(0, 2, 1)
print("source shape and strides", x.shape, x.strides)
print("transpose shape and strides", t.shape, t.strides)
print("transpose shares memory?", np.shares_memory(t, x))
print("transpose contiguous?", t.flags["C_CONTIGUOUS"])
reshaped = t.reshape(2, 12)
print("reshape shares memory with transpose?", np.shares_memory(reshaped, t))
print("reshaped shape", reshaped.shape)
print("reshaped snippet 0", reshaped[0])
assert np.array_equal(reshaped[0], [0, 4, 8, 1, 5, 9, 2, 6, 10, 3, 7, 11])

source shape and strides (2, 3, 4) (48, 16, 4)
transpose shape and strides (2, 4, 3) (48, 4, 16)
transpose shares memory? True
transpose contiguous? False
reshape shares memory with transpose? False
reshaped shape (2, 12)
reshaped snippet 0 [ 0.  4.  8.  1.  5.  9.  2.  6. 10.  3.  7. 11.]

Prefer reshape only when flat order is already correct. If downstream code requires a C-contiguous buffer, make that requirement explicit with np.ascontiguousarray(t) before reshaping. ravel() also returns a contiguous one-dimensional array, copying when needed, but flattening first discards the axis structure you were reasoning about.

Notice that the flattened snippet begins [0, 4, 8], not [0, 1, 2]. The default reshape follows C index order in the transposed array, not the physical order of its original buffer. That's why this operation must copy. A contiguous copy preserves the transposed values and axes; it doesn't undo the transpose.^{[4]Reference 4numpy.reshapehttps://numpy.org/doc/stable/reference/generated/numpy.reshape.html}

Axis names are positions, not portable labels. Image libraries may expect either (batch, channels, height, width) or (batch, height, width, channels), so write the intended axis order explicitly.

Matrix multiplication: compare every pair of tokens

So far, addition has paired features by position and mean has combined token rows. Another useful operation compares one token vector with another. Their dot product multiplies corresponding features and adds the products. Token the has vector [0, 1, 2, 3], and token red has vector [4, 5, 6, 7]; their dot product is $0\times4 + 1\times5 + 2\times6 + 3\times7 = 38$.

To compare every token with every other token in a snippet, we need nine such scores. Matrix multiplication, written @ in Python, computes those row-against-column dot products together. Unlike elementwise *, it combines a shared axis: (3, 4) @ (4, 3) gives (3, 3).^{[6]Reference 6numpy.matmulhttps://numpy.org/doc/stable/reference/generated/numpy.matmul.html}

This pairwise comparison is also part of attention, the operation that lets tokens draw information from other positions. A Transformer model builds a query vector for the token doing the comparison and a key vector for the token being compared.^{[7]Reference 7Attention Is All You Need.https://arxiv.org/abs/1706.03762} We can practice its array operations without learning the whole model yet.

Starting from our (B, T, D) array, multiply each token vector by learned matrices of shape (D, D) to produce queries and keys. For this exercise, use identity matrices: diagonal entries are 1 and the others are 0, so multiplying by them leaves the vectors unchanged.

q.shape == (B, T, D)
k.shape == (B, T, D)

Every query token should compare with T key tokens, so the score tensor should be (B, T, T). For arrays with more than two axes, @ treats the final two axes as matrices and broadcasts the earlier axes as a stack. Here that keeps the two snippets separate.^{[6]Reference 6numpy.matmulhttps://numpy.org/doc/stable/reference/generated/numpy.matmul.html}

The last two axes of k must swap first. Each (T, D) query matrix then multiplies a (D, T) key matrix:

unscaled_scores = q @ np.swapaxes(k, -1, -2)
unscaled_scores.shape == (B, T, T)

k has to change axis order before that multiply. q keeps (B, T, D) and goes straight in.

View larger

The inner size that has to agree is D. After the swap, that shared D disappears into the scores, leaving two T axes: which token asked and which token was compared.

Check the attention scores

Use identity projection matrices so q and k keep the same values as x. That lets you verify the first score by hand: $0^2 + 1^2 + 2^2 + 3^2 = 14$.

import numpy as np

B, T, D = 2, 3, 4
x = np.arange(B * T * D, dtype=np.float32).reshape(B, T, D)
wq = np.eye(D, dtype=np.float32)
wk = np.eye(D, dtype=np.float32)

q = x @ wq
k = x @ wk

try:
    q @ k
except ValueError:
    print("q @ k fails: inner axes are D=4 and T=3")
else:
    raise RuntimeError("Expected q @ k to fail for shapes (2, 3, 4) and (2, 3, 4)")

unscaled_scores = q @ np.swapaxes(k, -1, -2)

print("q", q.shape)
print("k", k.shape)
print("unscaled_scores", unscaled_scores.shape)
print("first query against first key", unscaled_scores[0, 0, 0])
print("token 'the' against token 'red'", unscaled_scores[0, 0, 1])
assert unscaled_scores[0, 0, 1] == 38

q @ k fails: inner axes are D=4 and T=3
q (2, 3, 4)
k (2, 3, 4)
unscaled_scores (2, 3, 3)
first query against first key 14.0
token 'the' against token 'red' 38.0

These are raw comparison scores, not probabilities or the final attention output. Later lessons turn them into weights and use those weights to combine token information. For now, the checkable contract is enough: scores[b, i, j] compares token i with token j in snippet b, using all four features.

Split each token's features, then move the new axis

Transformers often perform several attention comparisons in parallel, called heads.^{[7]Reference 7Attention Is All You Need.https://arxiv.org/abs/1706.03762} For a shape-only preview, split the four features of token the into two groups: [0, 1] and [2, 3]. With H equally sized heads, D must be divisible by H; each group has width D // H.

Splitting the last axis exposes (B, T, H, D // H). That still keeps all groups for a token together.

To finish at (B, H, T, D // H), first reshape to expose the head axis, then transpose to move it next to the batch axis. Reshaping straight to the target shape produces a legal array whose numbers are attached to the wrong head and token positions.

import numpy as np

B, T, D, H = 2, 3, 4, 2
assert D % H == 0
head_dim = D // H
x = np.arange(B * T * D, dtype=np.float32).reshape(B, T, D)

# Right: expose the head axis, then move it next to the batch axis.
heads = x.reshape(B, T, H, head_dim).transpose(0, 2, 1, 3)

# Wrong: same target shape, but the values land in the wrong heads.
jumbled = x.reshape(B, H, T, head_dim)

print("heads shape", heads.shape)
print("same positions?", np.array_equal(heads, jumbled))
print("head 0, token 'red'", heads[0, 0, 1])
print("direct reshape, same index", jumbled[0, 0, 1])
assert np.array_equal(heads[0, 0, 1], x[0, 1, :head_dim])

heads shape (2, 2, 3, 2)
same positions? False
head 0, token 'red' [4. 5.]
direct reshape, same index [2. 3.]

Both arrays have shape (2, 2, 3, 2), so a shape check alone would pass. But head 0 for token red must contain [4, 5], not the last two features of token the. The value assertion catches what the shape assertion can't.

Transfer: the same rule on images

The trailing-axis rule isn't special to text. A computer-vision pipeline might prepare 32 satellite tiles for a land-cover model. Each red-green-blue (RGB) tile is 224 pixels high by 224 pixels wide:

images = (32, 224, 224, 3)

Color lives on the last axis. Suppose pixel values are scaled between 0 and 1, and you want to subtract one training-set mean per color channel from every pixel. The example uses made-up channel means [0.52, 0.48, 0.44] and all-white tiles so the result is easy to inspect.

import numpy as np

images = np.ones((32, 224, 224, 3), dtype=np.float32)  # placeholder normalized tiles
mean_color = np.array([0.52, 0.48, 0.44], dtype=np.float32)  # illustrative means

normalized = images - mean_color
print("normalized shape", normalized.shape)
print("one centered pixel", normalized[0, 0, 0].round(2))

normalized shape (32, 224, 224, 3)
one centered pixel [0.48 0.52 0.56]

NumPy lines the shapes up from the right. The trailing 3 matches, so the mean color broadcasts across all 224 rows, 224 columns, and 32 photos. No Python loop is needed.

Common shape bugs: symptoms and fixes

Shape bugs become easier to fix when you name the symptom, the cause, and the correction.

When a shape bug appears, write the intended axis meaning before changing code. Then check whether the operation preserved it.

Build a shape guard script

Now turn the running example into a small boundary check. Start with x = np.arange(24).reshape(2, 3, 4), add a feature bias and keep (2, 3, 4), pool tokens and keep (2, 4), then build attention scores and keep (2, 3, 3). End by replacing the bias with shape (3,) and reading the rejection.

Wrap each operation in a function so its input and output contract stays visible:

• add_feature_bias(x, bias)
• pool_tokens(x)
• unscaled_attention_scores(x, wq, wk)

Try implementing these three functions before reading the solution. Reject a bad input shape explicitly and assert the output invariant each function promises. Python can remove assert statements when run with optimization enabled, so use ordinary exceptions for checks at a real data boundary. These functions check shapes only; add dtype checks separately if the caller requires a particular numeric type.

Run this reference version and compare its output with your predictions:

import numpy as np

def add_feature_bias(x: np.ndarray, bias: np.ndarray) -> np.ndarray:
    if x.ndim != 3:
        raise ValueError(f"expected (B, T, D), got {x.shape}")
    expected_bias_shape = (x.shape[-1],)
    if bias.shape != expected_bias_shape:
        raise ValueError(f"expected bias shape {expected_bias_shape}, got {bias.shape}")
    out = x + bias
    assert out.shape == x.shape
    return out

def pool_tokens(x: np.ndarray) -> np.ndarray:
    if x.ndim != 3:
        raise ValueError(f"expected (B, T, D), got {x.shape}")
    out = x.mean(axis=1)
    assert out.shape == (x.shape[0], x.shape[2])
    return out

def unscaled_attention_scores(x: np.ndarray, wq: np.ndarray, wk: np.ndarray) -> np.ndarray:
    if x.ndim != 3:
        raise ValueError(f"expected (B, T, D), got {x.shape}")
    batch, tokens, width = x.shape
    expected_weight_shape = (width, width)
    if wq.shape != expected_weight_shape or wk.shape != expected_weight_shape:
        raise ValueError(
            f"expected weight shapes {expected_weight_shape}, got {wq.shape} and {wk.shape}"
        )
    q = x @ wq
    k = x @ wk
    unscaled_scores = q @ np.swapaxes(k, -1, -2)
    assert unscaled_scores.shape == (batch, tokens, tokens)
    return unscaled_scores

x = np.arange(24, dtype=np.float32).reshape(2, 3, 4)
bias = np.array([100, 200, 300, 400], dtype=np.float32)
wq = np.eye(4, dtype=np.float32)
wk = np.eye(4, dtype=np.float32)

shifted = add_feature_bias(x, bias)
pooled = pool_tokens(shifted)
unscaled_scores = unscaled_attention_scores(x, wq, wk)

# Expected values come from the earlier feature sums and token dot products.
assert np.array_equal(shifted[0, 1], [104, 205, 306, 407])
assert np.array_equal(pooled[0], [104, 205, 306, 407])
assert np.array_equal(unscaled_scores[0], [
    [14, 38, 62],
    [38, 126, 214],
    [62, 214, 366],
])

print("shifted", shifted.shape)
print("pooled", pooled.shape)
print("unscaled_scores", unscaled_scores.shape)
print("pooled snippet 0", pooled[0])
print("score row for token 'the'", unscaled_scores[0, 0])

try:
    add_feature_bias(x, np.array([10, 20, 30], dtype=np.float32))
except ValueError as exc:
    print("caught:", exc)
else:
    raise RuntimeError("Expected the shape guard to reject a token-sized bias")

shifted (2, 3, 4)
pooled (2, 4)
unscaled_scores (2, 3, 3)
pooled snippet 0 [104. 205. 306. 407.]
score row for token 'the' [14. 38. 62.]
caught: expected bias shape (4,), got (3,)

When extending these checks, use both unequal axis sizes and equal ones. Unequal sizes exposed the earlier broadcasting error, but equal sizes can hide a misplaced axis. Pair shape checks with a known-value assertion, as the head-splitting example did for token red.

On a 24-number array, name the axis for snippets, the axis for tokens, and the axis for features. A reduction or reshape can keep the count and still drop the meaning. (B, T, D) records how many snippets, how many positions in each snippet, and how many numbers describe each position. Preserve those axis meanings alongside the element count. A useful check pairs the expected shape with one value you can trace back to its original token.