NumPy and Tensor Shapes

Two services emit short log snippets: "cache miss trace" and "timeout from worker." A routing model can't work directly with those words, so it receives small lists of numbers that represent each word. If a data-loading step swaps the service and word positions, later math can fail or, worse, run while attaching predictions to the wrong text.

Array shape is your first AI engineering habit. An array is an ordered collection of numbers; a multi-axis array is often called a tensor in machine learning. NumPy calls its array object an ndarray: it stores numerical data in a buffer plus metadata such as shape and strides that map each index to a position in that buffer. Views can reuse the same buffer with different metadata, so fast operations don't give you permission to ignore axis meaning.^{[1]Reference 1Copies and views.https://numpy.org/doc/stable/user/basics.copies.html}

That habit matters early in AI work. Many beginner bugs aren't mysterious math bugs. They are shape bugs:

• A batch axis disappears and the loss function crashes.
• A feature vector broadcasts across the wrong dimension and each token receives the wrong offset.
• A transpose should have happened, but a reshape happened instead and token meaning silently changes.

Use those two log snippets as one tiny batch. Afterward, name the axes in an array, predict how indexing or broadcasting changes its shape, reduce the right axis on purpose, choose reshape or transpose for the job, and check an attention-score tensor without guessing.

Your first batch: two log snippets

Use one tiny batch for the whole chapter.

Use two short log phrases:

A token is a piece of text a model processes; here we use one word as one token. A feature vector is the list of numbers representing one token. Later chapters call these per-position vectors hidden states. Suppose each token receives four features. The resulting tensor has shape:

Read that slowly:

That tuple is a contract, not decoration. If you can say the sentence, you can usually debug the code.

The anatomy of a tensor

Don't treat a tensor as a mysterious cube. Our tensor (2, 3, 4) is two snippet trays; each tray contains three token slots; each token slot contains four numerical features. To find one value, you need three addresses. Shape tells you how many choices exist at each address.

Keep the words concrete. If D means feature count, each vector on that axis should contain D numbers.

Indexing: which axes stay and which disappear

Start with real numbers so the array stops feeling abstract:

The shape is still (2, 3, 4), and each coordinate now points at an actual token vector.

The shape rule is simple:

1. Write the original shape.
2. A single number like 0 or 1 chooses one item and usually removes that axis.
3. A colon : keeps the whole axis.
4. Write the remaining axes in order.
5. Check with .shape.

Code lab: inspect one batch

Run this as a tiny script. Notice that it prints values as well as shapes.

The useful beginner habit isn't "look at .shape after the fact." It's "predict the shape, then let .shape confirm or correct you."

Broadcasting: reuse one feature bias

Now add one feature bias to each token vector:

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

One token by hand

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.^{[2]Reference 2Broadcasting.https://numpy.org/doc/stable/user/basics.broadcasting.html}

What broadcasting keeps

• the batch axis stays
• the token axis stays
• the feature bias repeats across both axes

The matching rule is easier than it sounds:

1. Line shapes up from the right.
2. Two axes are compatible if they are 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):

That last row is important. A shape can be wrong in two ways:

• it can fail loudly
• it can succeed but mean the wrong thing

The same rule as a shape-alignment diagram:

Code lab: broadcast one feature bias

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

That message is useful. It shows the feature axis was not where you thought it was.

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

The (N,) vs (N, 1) trap

A 1D array (5,) isn't the same as a column vector (5, 1). This is a common cause of silent broadcasting errors in neural networks.

A plain (3,) can broadcast into either a row or a column, depending on what it's paired with. That flexibility is convenient but dangerous. When you need an explicit row or column vector, make the axis visible with v[None, :] or v[:, None].

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:

For easier arithmetic, keep the same (B, T, D) contract but switch to feature values whose average you can check by hand:

Mean snippet 0 by hand

So x.mean(axis=1) should keep batch and feature axes, giving shape (B, D) = (2, 4).

What that reduction keeps

• the batch axis stays
• the token axis disappears
• the feature axis stays

Now compare that with:

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

Code lab: compare two reductions

To avoid reduction mistakes, say the axis name out loud before you write the code:

"I want one vector per snippet, so I am reducing the token axis."

Why keepdims=True is a safety switch

By default, a reduction removes the axis entirely. That can break later broadcasting because the shape has changed.

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.

This is easier to trust when you compare both shape paths side by side:

Reshape vs. transpose: regrouping or reordering

Beginners often mix these up because both can produce a new shape.

Use a smaller array so you can see the numbers move:

Both results have shape (2, 2, 3), but they answer different questions:

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.

One more beginner trap lives here:

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 vs. copies: who shares memory

Basic slicing produces a view: an array that reaches the same data buffer through different metadata. NumPy's docs state that basic indexing always creates views, while advanced indexing (integer arrays or boolean masks) always creates copies.^{[1]Reference 1Copies and views.https://numpy.org/doc/stable/user/basics.copies.html} transpose returns a view. reshape returns a view when the new layout can be expressed with strides, but it may copy a non-contiguous array. Those differences matter when you mutate arrays or reason about memory use.^{[1]Reference 1Copies and views.https://numpy.org/doc/stable/user/basics.copies.html}

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.

Contiguity after transpose: reshape may copy

A transpose keeps the same buffer but changes strides, so the result is often non-contiguous. reshape then may have to materialize a contiguous copy before regrouping. This commonly surprises head-split or head-merge code.

Prefer reshape only when flat order is already correct (or call .ravel() / make the array contiguous first when you need a guaranteed layout). Axis names are positions, not portable labels: libraries also disagree on layouts such as NCHW vs NHWC, so write the intended axis order explicitly.

From hidden states to attention scores

Attention looks harder than indexing or reduction, but the shape logic is the same. A later Transformer chapter will explain the full computation. For now, track one intermediate result: each token's query vector is compared with every token's key vector by a dot product before those scores become weights.^{[3]Reference 3Attention Is All You Need.https://arxiv.org/abs/1706.03762}

Start with hidden states of shape (B, T, D). Build queries and keys with matrices of shape (D, D). That keeps the batch axis and token axis in place:

To compare each token with each token, use matrix multiplication (many dot products calculated together). The last two axes of k must swap first:

Read the two T axes

• left T: which token is asking the question
• second T: which token is being compared against

Code lab: check attention scores

Get comfortable with shapes before Transformer internals. The score matrix is still built from small, explainable axis moves; later you'll scale it, apply a mask with shape (B, 1, T) or (B, T, T), run softmax, and use those weights to mix values. Those steps keep the score/weight tensor at shape (B, T, T) (or with a head axis when multi-head); matmul alone is not full attention.

Splitting heads: reshape then transpose

Transformers run several attention heads in parallel by carving feature axis D into H smaller feature groups of width D // H.^{[3]Reference 3Attention Is All You Need.https://arxiv.org/abs/1706.03762} Same reshape-versus-transpose trap as earlier. To go from (B, T, D) to (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 jumbled.

Both arrays have shape (2, 2, 3, 4), so a shape check alone would pass. Only the reshape-then-transpose path keeps each token's features grouped inside the right head.

Transfer: normalizing satellite tiles

The same rule applies outside text. Suppose a computer-vision pipeline prepares 32 satellite tiles for a land-cover model. Each red-green-blue (RGB) tile is 224 pixels high by 224 pixels wide. The shape is:

The last axis is the color channel. Pixel values are scaled between 0 and 1, and you want to subtract a measured mean-color vector (3,) from each pixel in every photo.

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.

Every fix here starts the same way: write the intended axis meaning, then check whether the operation preserves it.

Shape contract: The error message operands could not be broadcast together isn't a cryptic NumPy complaint. It's a precise statement that two axes you assumed matched don't match. Treat it as a shape-contract violation, not a random failure, and the next check becomes obvious.

Build a shape guard script

Build one tiny script that proves you understand the whole chapter.

1. Create x = np.arange(24).reshape(2, 3, 4).
2. Add a feature bias of shape (4,) and assert the shape stays (2, 3, 4).
3. Average over the token axis and assert the shape becomes (2, 4).
4. Create unscaled attention scores and assert the shape becomes (2, 3, 3).
5. Intentionally replace the bias with shape (3,) and read the failure.

If you want a stronger version, wrap each step in a function:

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

Each function should assert both the input shape it expects and the output shape it promises.

Use this compact reference version:

Check your reasoning: predict before you run

Cover the code blocks above and answer these without running anything. If your explanation names numbers without axis meanings, slow down and add labels: reliable shape reasoning attaches each number to its role.

Open each checkpoint only after you commit to an answer.

You have a batch of 8 log snippets, each 16 tokens long, with 512-dimensional embeddings. What is the shape? What does x[:, 0, :] return?

Answer

The full tensor shape is (8, 16, 512). x[:, 0, :] keeps every batch item, picks token position 0, and keeps the full feature axis, so the result shape is (8, 512).

You add a feature bias of shape (512,) to that batch. Does broadcasting succeed? What is the output shape?

Answer

Yes. Broadcasting succeeds because NumPy lines shapes up from the right, and the trailing 512 matches the feature axis. The output shape stays (8, 16, 512).

You average over the token axis to get one vector per snippet. What shape results? What shape results if you use keepdims=True?

Answer

Reducing the token axis means axis=1, so x.mean(axis=1) gives one feature vector per snippet with shape (8, 512). If you keep the reduced axis as a size-1 placeholder, x.mean(axis=1, keepdims=True) gives (8, 1, 512).

You want to turn (8, 16, 512) into (8, 512, 16) so the feature axis sits next to the batch axis. Should you use reshape or transpose?

Answer

Use transpose, because you are permuting axis meaning rather than regrouping the flat data stream. reshape can give the same numbers arranged into a legal new shape while silently changing what each position means.

You have queries of shape (8, 16, 512) and want 8 attention heads. What is the target shape, and what two operations get you there in the right order?

Answer

With 8 heads, each head is 512 / 8 = 64 wide, so the target is (8, 8, 16, 64) as (B, H, T, head_dim). First reshape to (8, 16, 8, 64) to expose the head axis, then transpose(0, 2, 1, 3) to move heads next to the batch axis. Reshaping straight to (8, 8, 16, 64) gives the same shape but jumbles which features belong to which head.

If you got any of these wrong, re-read the one-running-example section and trace the indices by hand before moving on.

Keep the axis contract visible

A shape is a sentence about your data, not a trio of integers.

(B, T, D) tells you how many examples you have, how many positions sit in each example, and how many numbers describe each position. Keep that sentence attached to the array and matrix multiplies, reductions, and attention stay readable instead of opaque.