NumPy and Tensor Shapes

Two customers send support tickets: "refund request order" and "shipped from hub." 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 customer and word positions, later math can fail or, worse, run while attaching predictions to the wrong text.

This is why 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 holds numerical data in a buffer plus metadata such as shape and strides that tell NumPy how an index reaches that data. Views can reuse the same buffer with different metadata, so fast operations aren't permission to ignore axis meaning.^[1]

That habit matters early in AI work. Many beginner bugs are not 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.

This chapter uses those two support tickets as one tiny batch. By the end, you should be able to 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 support tickets

Use one tiny batch for the whole chapter.

Imagine two short support-ticket 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 ticket 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), but now you can point at actual token vectors.

The pattern 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 ticket without a Python loop. You can imagine the bias vector repeating across the batch and token axes, but the repetition is conceptual: broadcasting doesn't need to copy that smaller operand in memory.^[2]

What broadcasting keeps

• the batch axis stays
• the token axis stays
• the feature bias repeats across both of those 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. 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

Here is the same rule shown 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 tells you 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,) is not 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 ticket, 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 ticket 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

The easiest way to avoid reduction mistakes is to say the axis name out loud before you write the code:

"I want one vector per ticket, 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 is 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. transpose also returns a view whenever possible. reshape returns a view when the new layout can be expressed with strides, but it can make a copy for a non-contiguous input. Fancy indexing (an integer array or a boolean mask) returns a copy. Those differences matter when you mutate arrays or reason about memory use.^[1]

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.

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]

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

This is why shape practice belongs before Transformer internals. The score matrix is still built from small, explainable axis moves; later you will scale it, hide positions a token shouldn't read, turn scores into weights, and use those weights to mix values.

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] This is exactly the reshape-versus-transpose trap from 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 batch axis. Reshaping straight to 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 catalog photos

The same rule applies outside ticket text. Suppose ShopFlow prepares 32 catalog photos for a product-image model. Each red-green-blue (RGB) photo 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.

The pattern behind all of these fixes is the same: write the intended axis meaning before asking whether the operation preserves that meaning.

Key insight: 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 result is still (2, 3, 4).
3. Average over the token axis and assert the result is (2, 4).
4. Create unscaled attention scores and assert the result is (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.

Here is a 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.

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

The durable lesson is simple: a shape is a sentence about data.

(B, T, D) is not three numbers. It says how many examples you have, how many positions each example contains, and how many numbers describe each position. Once you keep that sentence attached to the array, matrix multiplies, reductions, and attention stop feeling opaque.

Mastery check

Key concepts

• ndarrays
• axis naming
• indexing
• broadcasting
• axis reductions
• keepdims
• reshape vs transpose
• matrix multiplication
• shape assertions
• column vs row vectors

Evaluation rubric

• Foundational: Reads tensor axes in plain English before slicing, broadcasting, or reducing arrays
• Intermediate: Predicts and verifies shapes for indexing, broadcasting, reductions, projection, and attention-score examples
• Intermediate: Uses keepdims=True to preserve axis positions during reductions and avoid silent broadcasting mismatches
• Advanced: Adds shape assertions that catch wrong broadcasting, wrong reduction axes, reshape-versus-transpose confusion, and dropped batch dimensions

Follow-up questions

Common pitfalls

"Broadcast succeeded, so shape must be right"

Symptom: Code runs and output shape looks plausible, but every token or package gets the wrong offset. Cause: A legal broadcast matched the wrong axis. NumPy checked size compatibility, not semantic intent. Fix: Name the target axis first, then verify the smaller array is shaped for that axis on purpose.

"Reduction returned numbers, so averaging step was correct"

Symptom: mean() runs without complaint, but downstream code gets scalars where vectors were expected. Cause: The wrong axis was reduced, so the result kept different meaning than intended. Fix: Say the axis name out loud before reducing, then check kept axes against the plain-English goal.

"Reshape and transpose are interchangeable"

Symptom: Final shape matches expectation, but token order, catalog-photo layout, or attention heads become mixed up. Cause: reshape regrouped the flat stream when the real task was axis permutation. Fix: Use transpose when axis meaning moves. Use reshape only when flat order is already correct.

"A 1D vector is close enough to row or column"

Symptom: (N,) sometimes behaves like a row and sometimes like a column, creating silent broadcast surprises. Cause: One axis is missing entirely, so NumPy infers placement from context. Fix: Make the axis explicit with v[None, :] or v[:, None] whenever orientation matters.

"Slice is harmless to mutate"

Symptom: Editing a slice unexpectedly changes the original batch. Cause: Basic slices are views, and reshape may return a view that shares memory with the source. Fix: Call .copy() when you need an independent array before mutating it.