13. NumPy: the array model everything else copies

A NumPy array is one contiguous block of raw numbers plus metadata (dtype, shape, strides) describing how to read it. PyTorch tensors and JAX arrays use the same model, so learning it here pays off across the whole AI stack.

Anatomy of an array

import numpy as np

x = np.arange(12, dtype=np.float32).reshape(3, 4)
x.dtype      # float32: every element has the same type, 4 bytes each
x.shape      # (3, 4)
x.strides    # (16, 4): bytes to jump to the next row, and to the next column
x.nbytes     # 48
x.flags["C_CONTIGUOUS"]   # True: stored row by row

• One dtype for all elements means no per-element objects and no pointers. The CPU streams straight through memory.
• Strides make many operations free. A transpose just swaps the strides, and a slice changes the starting offset and strides. No data moves.

Views and copies

• Views share memory with the original: basic slicing (x[1:, ::2]), .T, and reshape when the layout allows it.
• Copies are new arrays: fancy indexing (x[[0, 2]]), boolean masks (x[x > 3]) and arithmetic results.

row = x[0]                 # a view
row[:] = 0                 # this modifies x too!
safe = x[0].copy()
np.shares_memory(row, x)   # True

Non-contiguous views, such as a transpose, can slow later operations or force hidden copies (np.ascontiguousarray fixes the layout). PyTorch has the same distinction: .view() requires contiguous memory, .reshape() copies when it must, and .contiguous() forces a contiguous copy.

Vectorization and reductions

Universal functions (np.exp, +, np.maximum) apply elementwise in compiled code. Reductions such as sum, max and mean take an axis:

logits = np.random.default_rng(0).normal(size=(32, 10))     # (batch, classes)
z = logits - logits.max(axis=1, keepdims=True)              # subtract the max for stability
probs = np.exp(z) / np.exp(z).sum(axis=1, keepdims=True)    # softmax per row

axis=k means “collapse dimension k”. keepdims=True keeps it as size 1, so the result still lines up against the input.

Broadcasting

Broadcasting lets arrays of different shapes combine without copying. Compare shapes from the right:

1. Two dimensions are compatible if they are equal, or if one of them is 1.
2. Missing leading dimensions count as 1.
3. Size-1 dimensions are stretched virtually, with stride 0 and no copy.

(32, 10)    - (32, 1)      -> (32, 10)       per-row max: OK
(32, 10)    + (10,)        -> (32, 10)       per-class bias: OK
(32, 10)    + (32,)        -> error          aligned right, 10 vs 32
(32, 1, 64) - (1, 50, 64)  -> (32, 50, 64)   all pairwise differences: OK

The danger is that broadcasting silently succeeds on the wrong shapes. With y_true of shape (n,) and y_pred of shape (n, 1), y_true - y_pred has shape (n, n). The loss looks fine and is wrong, so assert shapes.

None (or np.newaxis) inserts a size-1 axis: a[:, None] - b[None, :] builds a pairwise matrix.

Linear algebra

• a @ b is matrix multiplication, batched over leading dimensions for 3D and higher arrays.
• np.einsum("bij,bjk->bik", a, b) names every axis explicitly; it is excellent for reading and writing attention code.
• np.linalg.norm(v, axis=-1, keepdims=True) normalizes vectors. After normalizing embeddings, cosine similarity is just a matrix product.

Brute-force semantic search, the core of a vector database before any indexing tricks:

def cosine_topk(query: np.ndarray, docs: np.ndarray, k: int = 5) -> np.ndarray:
    # query: (d,), docs: (n, d)
    q = query / np.linalg.norm(query)
    d = docs / np.linalg.norm(docs, axis=1, keepdims=True)
    scores = d @ q                              # (n,)
    idx = np.argpartition(-scores, k)[:k]       # top k in O(n), unsorted
    return idx[np.argsort(-scores[idx])]        # sort only those k

dtypes and randomness

• NumPy defaults to float64; machine learning usually wants float32. Mixing types silently promotes to the wider one.
• Integer arrays overflow and wrap around, unlike Python ints.
• Use a local generator for randomness: rng = np.random.default_rng(seed), then rng.normal(...). It is reproducible and avoids hidden global state.

Questions to ask:

• Is this a view or a copy?
• What is the shape after each line, and is broadcasting doing what I think?
• Which axis am I reducing over?
• What dtype is this, and is something silently promoting to float64?