NumPy and Tensor Shapes

Two log snippets contain three words each: "cache miss trace" and "timeout from worker." 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:

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 miss and returns its four-feature vector.

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 without a dtype uses NumPy's default integer. 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.

Selecting miss 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?

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

For the miss token, the four additions are:

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 usually avoids copying the smaller operand in memory.^{[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.

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).

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

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:

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:

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:

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).

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

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

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:

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.

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.

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)?

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.

Reshape vs. transpose: regrouping or reordering

At this point, x has named axes. The next operation asks for shape (2, 4, 3). Before running it, decide whether that means regrouping the data stream or moving the existing axes.

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

Both results have shape (2, 4, 3), but they answer different questions. Predict where the highlighted 3 should land if it remains cache's last feature. reshape parks it next to 4 (the first feature of miss), while transpose keeps it with cache in a feature-by-token grid.

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.^{[3]Reference 3numpy.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,).

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 stores a buffer of values and metadata describing how to reach them. Its strides say how many bytes to move when an index increases by one along an axis. 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.^{[4]Reference 4Copies 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.^{[4]Reference 4Copies 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.

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.

A transpose can make the next reshape copy

An array is C-contiguous when its values occupy one uninterrupted memory region in last-axis-first order: move across a row, then to the next row. A transpose changes how indices traverse that 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 sharing before and after that reshape:

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.

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. For cache = [0, 1, 2, 3] and miss = [4, 5, 6, 7], that gives $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).^{[5]Reference 5numpy.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.^{[6]Reference 6Attention 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.

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.^{[5]Reference 5numpy.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:

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

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$.

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.^{[6]Reference 6Attention Is All You Need.https://arxiv.org/abs/1706.03762} For a shape-only preview, split each token's four features into two groups: cache becomes [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.

Both arrays have shape (2, 2, 3, 2), so a shape check alone would pass. But head 0 for miss must contain [4, 5], not the last two features of cache. 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:

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.

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:

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 the miss token.

Your loader now returns one snippet as shape (3, 4). The shape guard rejects it. How can you add a batch axis without changing which values belong to each token?

Answer

For single.shape == (3, 4), use single[None, :, :] to get (1, 3, 4). The new leading axis says there's one snippet. If you're selecting from the original batch instead, x[0:1] preserves that leading axis directly, whereas x[0] removes it.

(B, T, D) records how many snippets, how many positions in each snippet, and how many numbers describe each position. Preserve those meanings, not just the element count. A useful check pairs the expected shape with one value you can trace back to its original token.