NumPy and Tensor Shapes

NumPy teaches you to see data as shapes.

Python gave you rows and functions. NumPy adds a new rule: every value has axes, and those axes must mean something.

In LLM work, many bugs look like mysterious math bugs. Often they aren't mysterious. Batch, token, and hidden dimensions got swapped, squeezed, or broadcast in the wrong place.

What You Learn

This chapter teaches the shape habit. By the end, a beginner should be able to look at an array, name each axis, predict the output shape of a small operation, and write assertions that catch wrong shapes before model code runs.^[1][2][3]

Step Map

Analogy: a tensor is like a set of labeled shelves. If you remove the labels, the boxes may still stack, but nobody knows what the stack means.

---

Tiny Story

Imagine two short prompts:

Each token becomes four numbers.

That means the hidden states have shape:

text
1(B, T, D) = (2, 3, 4)

Read it slowly:

The shape isn't decoration. It is the contract.

Vocabulary

• array: grid of numbers.
• axis: one direction in an array.
• shape: tuple of axis sizes, such as (2, 3, 4).
• scalar: one number with shape ().
• vector: one axis, such as (4,).
• matrix: two axes, such as (3, 4).
• tensor: array with any number of axes.
• broadcasting: NumPy rule that stretches compatible axes.

Keep the words concrete. If D means hidden size, every vector on that axis should have D numbers.

Worked Example: Indexing

Start with shape (2, 3, 4).

Indexing is reading from the shelf.

When you remove one coordinate, you usually remove one axis.

Code Lab

Put this in numpy_shapes_demo.py:

python
1import numpy as np
2
3B, T, D = 2, 3, 4
4x = np.arange(B * T * D).reshape(B, T, D)
5
6print("x", x.shape)
7print("first_prompt", x[0].shape)
8print("one_token", x[0, 1].shape)
9print("all_second_tokens", x[:, 1].shape)
10
11assert x.shape == (2, 3, 4)
12assert x[0, 1].shape == (4,)

Expected output:

text
1x (2, 3, 4)
2first_prompt (3, 4)
3one_token (4,)
4all_second_tokens (2, 4)

The values inside x aren't important yet. The shape changes are the lesson.

Second Worked Example: Matrix Multiply

Now transform each token vector from four numbers to five numbers.

Use a weight matrix:

The last axis of x must match the first axis of W.

Shape rule:

text
1(2, 3, 4) @ (4, 5) -> (2, 3, 5)

The 4 disappears because it is the matched dimension. The 5 appears because it is the new feature size.

Run The Projection

python
1import numpy as np
2
3B, T, D, H = 2, 3, 4, 5
4x = np.ones((B, T, D))
5W = np.ones((D, H))
6y = x @ W
7
8print("y", y.shape)
9assert y.shape == (B, T, H)

Expected output:

text
1y (2, 3, 5)

That one rule appears everywhere in neural networks.

Third Worked Example: Attention Scores

Attention compares tokens to tokens.

For one batch item with three tokens, each token needs a score against three tokens. That gives a T x T table.

For a full batch, the score shape should be:

text
1(B, T, T)

Shape Checkpoint

Build the smallest version:

python
1import numpy as np
2
3B, T, D = 2, 3, 4
4rng = np.random.default_rng(7)
5x = rng.normal(size=(B, T, D))
6wq = rng.normal(size=(D, D))
7wk = rng.normal(size=(D, D))
8
9q = x @ wq
10k = x @ wk
11scores = q @ np.swapaxes(k, -1, -2)
12
13print("q", q.shape)
14print("k", k.shape)
15print("scores", scores.shape)
16assert scores.shape == (B, T, T)

Expected output shape:

text
1scores (2, 3, 3)

Read The Output

This is the first shape checkpoint for Transformer attention. You don't need full attention yet. You only need to know why token-to-token scores have two token axes.

Common Trap

The common trap is trusting NumPy because it returned an array.

NumPy can return a legal array that means the wrong thing.

Example:

The fix isn't more confidence. The fix is shape assertions with names.

Mini Test

Add tests like these:

python
1import numpy as np
2
3def test_token_projection_shape():
4    B, T, D, H = 2, 3, 4, 5
5    x = np.ones((B, T, D))
6    W = np.ones((D, H))
7    y = x @ W
8    assert y.shape == (B, T, H)
9
10def test_attention_scores_are_token_by_token():
11    B, T, D = 2, 3, 4
12    q = np.ones((B, T, D))
13    k = np.ones((B, T, D))
14    scores = q @ np.swapaxes(k, -1, -2)
15    assert scores.shape == (B, T, T)

These tests don't prove a model is good. They prove the arrays still mean what their names claim.

Practice

1. Change T from 3 to 6.
2. Predict every printed shape before running code.
3. Change H from 5 to 2.
4. Explain why y changes only on the last axis.
5. Remove np.swapaxes and read the error.
6. Write one sentence that starts: "The scores shape is (B, T, T) because..."

Production Check

At every model boundary, write down:

• input shape
• axis names
• expected output shape
• batch behavior
• sequence behavior
• feature or hidden size
• shape assertion
• example with tiny numbers

For example:

Input hidden states are (B, T, D). Projection matrix is (D, H). Output is (B, T, H). Test uses B=2, T=3, D=4, H=5.

That note prevents the most common beginner mistake: treating arrays as anonymous grids.

Next, continue to Data Structures for AI. NumPy taught you to keep shape meaning attached to arrays. Data structures teach you to keep lookup meaning attached to storage.