NumPy: The Array That Everything Else Is Built On

Pandas is NumPy with labels. Scikit-learn's estimators take and return NumPy arrays. PyTorch tensors deliberately copy NumPy's API and can share memory with it. Every plotting library consumes it. If you understand NumPy's memory model properly, most of the "why is this slow?" questions in the rest of the stack answer themselves.

Why NumPy is fast (and Python is not)

A Python list of a million integers is a million pointers to a million boxed PyObjects scattered across the heap. Each has a reference count, a type pointer, and its value. Iterating it means chasing pointers, checking types, and allocating result objects.

A NumPy array of a million int64s is one contiguous 8 MB block plus a small header. Operating on it dispatches once to a C loop that the CPU can vectorise with SIMD instructions and prefetch perfectly.

The measured gap is 10–100×, and it comes from three things: no per-element Python object overhead, cache-friendly contiguous access, and SIMD.

import numpy as np, time

a = list(range(1_000_000))
b = np.arange(1_000_000)

t = time.perf_counter(); s = sum(x * x for x in a); py = time.perf_counter() - t
t = time.perf_counter(); s = (b * b).sum();         np_ = time.perf_counter() - t
print(f"python {py*1000:.1f} ms   numpy {np_*1000:.1f} ms   speedup {py/np_:.0f}x")

The ndarray: four pieces of metadata

Every array is a buffer plus four descriptors. Understanding them explains reshape, transpose, slicing, and every "why did this copy?" question.

a = np.arange(12).reshape(3, 4)
a.dtype     # dtype('int64')
a.shape     # (3, 4)
a.strides   # (32, 8)  -> 4 elements * 8 bytes per row, 8 bytes per column
a.flags     # C_CONTIGUOUS: True, OWNDATA: False (it's a view of arange's buffer)

Strides are the key idea. A transpose does not move any data — it swaps the strides:

a.T.strides      # (8, 32)  -- same buffer, different reading order
a.T.flags        # F_CONTIGUOUS: True, C_CONTIGUOUS: False

That is why a.T is \(O(1)\) and why a.T.reshape(-1) is not — reshaping a non-contiguous array may need a copy, depending on stride compatibility. For example, np.arange(12)[::2].reshape(2, 3) still shares memory. The displayed transpose-flatten operation copies in its default C order. np.ascontiguousarray copies only when needed; .ravel() may copy and .flatten() always copies.

dtypes, and the memory they cost

Two defaults that cost real money. NumPy defaults to float64, and scikit-learn/PyTorch mostly want float32. A 10M-row × 100-feature matrix is 8 GB in float64 and 4 GB in float32. Always be explicit:

X = np.asarray(data, dtype=np.float32)

And dtype=object arrays give you none of NumPy's benefits — they store pointers to Python objects and every operation falls back to the interpreter. Some pandas string columns use object storage, but pandas 3 infers a dedicated string dtype and may use Arrow-backed storage. Storage and operation determine performance; pandas strings are not universally arrays of Python objects.

Integer overflow is silent. NumPy does not promote to bignum:

np.array([2**62], dtype=np.int64) * 2   # [-9223372036854775808]: signed overflow

Broadcasting

Broadcasting lets arrays of different shapes combine without materialising copies. The rules, applied right-to-left across the shapes:

1. If the arrays have different numbers of dimensions, left-pad the shorter with 1s.
2. Two dimensions are compatible if they are equal, or if one of them is 1.
3. A dimension of size 1 is stretched (by setting its stride to 0 — no data is duplicated).
4. Any other mismatch is an error.

X = np.random.randn(1000, 20)

# standardise every column, no loops, no temporaries you did not intend
Xz = (X - X.mean(0)) / X.std(0)          # (1000,20) - (20,) -> broadcast

# all pairwise squared distances between two point sets
A = np.random.randn(500, 3)
B = np.random.randn(800, 3)
D2 = ((A[:, None, :] - B[None, :, :]) ** 2).sum(-1)     # (500, 800)

The trap in that last example: A[:, None, :] - B[None, :, :] materialises a (500, 800, 3) intermediate: 9.6 MB here, but 60 GB for 50k by 50k float64 three-dimensional points. The identity avoids the three-dimensional temporary, not the 20 GB quadratic output or every other temporary:

D2 = (A**2).sum(1)[:, None] - 2 * A @ B.T + (B**2).sum(1)[None, :]
np.maximum(D2, 0, out=D2)     # clamp tiny negatives from float cancellation

That is the trick behind sklearn.metrics.pairwise_distances, and the clamp is there because catastrophic cancellation can produce small negative "squared" distances.

np.newaxis (a.k.a. None) is how you control which axis broadcasts. When a broadcast fails, print the shapes — the error message tells you the axis, and the fix is almost always an inserted None.

Vectorization: replacing loops with array expressions

The rule is simple: if you wrote a Python for loop over array elements, there is a faster way.

# a decision rule, three ways
scores = np.random.randn(1_000_000)

labels = np.array([1 if s > 0.5 else (-1 if s < -0.5 else 0) for s in scores])  # ~400 ms
labels = np.where(scores > 0.5, 1, np.where(scores < -0.5, -1, 0))              # ~6 ms
labels = np.sign(scores) * (np.abs(scores) > 0.5)                               # ~4 ms

np.vectorize is not vectorisation. It is a convenience wrapper around a Python loop and is roughly as slow. If the docs for a function say "provided primarily for convenience, not for performance", believe them.

Ufuncs and their extras

Universal functions (np.add, np.exp, np.maximum, …) all support the same machinery:

np.add(a, b, out=c)              # write into an existing buffer, no allocation
np.add.reduce(a, axis=0)         # == a.sum(0)
np.add.accumulate(a)             # == np.cumsum(a)
np.add.outer(a, b)               # all pairwise sums
np.add.at(a, idx, vals)          # UNBUFFERED add — handles repeated indices

np.add.at deserves emphasis. a[idx] += vals with repeated indices in idx only applies the last one, because the fancy-index read and write are separate operations. np.add.at(a, idx, vals) accumulates correctly. This is exactly the scatter-add that embedding-gradient accumulation needs, and it is a classic silent-wrong-answer bug.

Indexing: views vs copies

This single distinction causes more NumPy bugs than anything else.

a = np.arange(10)
v = a[2:5]        # view
v[0] = 999
a                 # array([0, 1, 999, 3, 4, 5, ...])  <- a was modified!

m = a[a > 3]      # copy
m[0] = -1
a                 # unchanged

Use np.shares_memory(a, b) to check overlap with the particular source array. .base can point to an intermediate copy, so it does not prove aliasing with the original input. When a function takes an array it may mutate, take a defensive .copy().

Fancy indexing patterns worth knowing

# one-hot encoding without sklearn
onehot = np.eye(n_classes, dtype=np.float32)[labels]         # (N, C)

# gather the predicted probability of the true class for every row
p_true = probs[np.arange(len(labels)), labels]               # (N,)

# top-k indices per row, unsorted (O(n)) then sorted within the k
assert 0 <= k <= scores.shape[1]
idx = (np.argpartition(-scores, k - 1, axis=1)[:, :k] if k
       else np.empty((len(scores), 0), dtype=int))
idx = np.take_along_axis(idx, np.argsort(-np.take_along_axis(scores, idx, 1), 1), 1)

# shuffle features and labels together
perm = rng.permutation(len(X)); X, y = X[perm], y[perm]

np.argpartition is the one people miss: it finds the top-\(k\) in \(O(n)\) instead of \(O(n\log n)\), which matters when \(n\) is a million and \(k\) is 10.

Reductions and axes

The axis argument means "the axis that disappears".

a = np.arange(24).reshape(2, 3, 4)
a.sum(axis=0).shape      # (3, 4)   -- axis 0 collapsed
a.sum(axis=(0, 2)).shape # (3,)
a.sum(axis=-1).shape     # (2, 3)
a.sum(axis=1, keepdims=True).shape   # (2, 1, 4)  -- kept for broadcasting

keepdims=True is what makes normalisation expressions work without manual None insertion:

probs = e / e.sum(axis=-1, keepdims=True)

The ddof mismatch between NumPy (0) and pandas (1) silently changes your reported standard deviation. Be explicit if the number matters.

Arithmetic reductions such as mean, sum, and max commonly propagate NaN. This is not true of every reduction: any/all treat NaN as truthy, and arg-reductions return indices. Define a finite-value policy; nanmean on an all-NaN slice remains undefined and emits a warning. Find invalid values first:

np.isnan(X).any(0)          # which columns contain NaN
np.isfinite(X).all()        # any NaN or inf anywhere?

Linear algebra

A @ B                       # matmul — use this, not np.dot, for 2-D
np.linalg.solve(A, b)       # solve Ax=b   -- NOT inv(A) @ b
np.linalg.lstsq(X, y, rcond=None)   # least squares, QR/SVD based
np.linalg.inv(A)            # you almost never want this
np.linalg.pinv(A)           # pseudo-inverse via SVD, handles singular A
np.linalg.eigh(S)           # symmetric eigendecomposition — faster and stabler than eig
np.linalg.svd(A, full_matrices=False)
np.linalg.cholesky(S)       # for positive-definite S; twice as fast as LU
np.linalg.norm(x, ord=2)    # ord: 1, 2, inf, 'fro', 'nuc'
np.linalg.matrix_rank(A)
np.trace(A), np.linalg.det(A), np.linalg.slogdet(A)

Three rules that come straight from numerical analysis:

1. Prefer solve to explicitly forming an inverse for a linear solve. It avoids unnecessary work and rounding operations; it does not improve the underlying system's conditioning.
2. Prefer lstsq to normal equations. Forming \(X^TX\) squares the spectral condition number for full-column-rank \(X\); inversion is not what causes that squaring.
3. Use slogdet instead of log(det(A)). Determinants of large matrices overflow or underflow; slogdet returns the sign and the log magnitude separately.

einsum: one notation for all of them

Einstein summation names the axes and sums over any index that does not appear in the output.

np.einsum('ij,jk->ik', A, B)          # matrix multiply
np.einsum('ij,ij->', A, B)            # Frobenius inner product
np.einsum('ii->', A)                  # trace
np.einsum('ij->ji', A)                # transpose
np.einsum('bij,bjk->bik', X, Y)       # batched matmul
np.einsum('bhqd,bhkd->bhqk', Q, K)    # attention scores, batch and head aware
np.einsum('i,j->ij', a, b)            # outer product
np.einsum('bi,bi->b', a, b)           # row-wise dot product

einsum is self-documenting in a way that a chain of transpose/reshape/ matmul is not, which matters enormously when reading attention code. Pass optimize=True for multi-operand contractions so it picks a good contraction order — the difference can be orders of magnitude.

Random numbers, done correctly

The legacy np.random.seed / np.random.rand global API is discouraged. Use the Generator API:

rng = np.random.default_rng(42)

rng.random((3, 4))                 # uniform [0,1)
rng.standard_normal((3, 4))        # N(0,1)
rng.normal(loc=0, scale=2, size=5)
rng.integers(0, 10, size=5)        # high exclusive, also true of np.random.randint
rng.choice(n, size=k, replace=False)
rng.permutation(n)
rng.shuffle(a)                     # in place

Generators are independent objects, so parallel workers each get their own stream: rng.spawn(n) constructs well-separated streams with extremely high probability, not a proof of mathematical independence. That is a useful correct pattern for dataloader workers, and it fixes the classic bug where every forked worker produces identical "random" augmentations.

Performance rules

# the shape of a memory bug
N, D = 100_000, 512
print(f"{N * D * 4 / 1e9:.1f} GB")        # 0.2 GB, fine
print(f"{N * N * 4 / 1e9:.1f} GB")        # 40.0 GB: an accidental pairwise matrix

np.memmap lets you work with arrays larger than RAM by mapping a file:

X = np.memmap('features.f32', dtype=np.float32, mode='r', shape=(10_000_000, 128))
batch = np.asarray(X[i:i+1024])       # only this slice is paged in

This is how large embedding tables and pretokenised datasets are usually stored.

Interoperability

import torch
t = torch.from_numpy(a)      # SHARES memory — mutating one mutates the other
b = t.numpy()                # also shares (CPU tensors only)
c = t.cpu().numpy().copy()   # explicit copy when you want independence

from_numpy is a specific shared-storage bridge, distinct from the DLPack APIs. It requires supported CPU dtypes and strides; negative-stride arrays need a copy. Do not mutate through a tensor backed by a read-only array. For tensors requiring gradients, explicitly detach before NumPy conversion when abandoning autograd is intended; a CUDA tensor must also move to the host.

Similarly, homogeneous pandas dtypes do not guarantee zero-copy to_numpy(): block layout, extension arrays, requested dtype and CoW matter. scipy.sparse matrices interoperate via .toarray() — which materialises the dense form, so check the size first.

Common bugs

Self-check

Worked CPU invariants

The fixture distinguishes ownership from overlap, handles constant-column scaling, and computes stable softmax. NumPy 2 promotion can preserve a low-precision array dtype when adding Python scalars; never infer accumulator precision from the literal alone. Use explicit dtypes and astype(casting="safe") when rejecting lossy conversion is the intended contract. The copies guide and promotion rules describe these version-sensitive semantics.

import numpy as np
a = np.arange(12)
strided = a[::2].reshape(2, 3)
assert np.shares_memory(a, strided)
assert not np.shares_memory(a, a.reshape(3, 4).T.reshape(-1))
assert np.isscalar(a[3]) and np.shares_memory(a, a[3:4])
counts = np.zeros(3, dtype=int)
np.add.at(counts, [1, 1, 2], 1)
assert counts.tolist() == [0, 2, 1]
assert (np.array([2**62], dtype=np.int64) * 4).item() == 0
assert 50_000 * 50_000 * 3 * 8 == 60_000_000_000
X = np.array([[1., 2.], [1., 4.]])
std = X.std(0)
normalized = np.divide(X-X.mean(0), std, out=np.zeros_like(X), where=std!=0)
assert np.allclose(normalized, [[0, -1], [0, 1]])
logits = np.array([[1000., 1001., 999.]])
exp = np.exp(logits-logits.max(-1, keepdims=True))
prob = exp/exp.sum(-1, keepdims=True)
assert np.isfinite(prob).all() and np.allclose(prob.sum(-1), 1)
print("aliasing, overflow, allocation, normalization, softmax checks passed")

Why is the out= initialization important? A ufunc's false where positions retain the existing output; an uninitialized output would contain arbitrary values. Why can Gram distances be inaccurate? Large nearly equal norm and dot-product terms cancel; clamping negative results does not recover lost significant digits. For nearest neighbors, process bounded blocks and retain only current top-k results instead of allocating the entire quadratic output. Benchmark repeated runs with matched dtypes, warmup, thread counts, and peak memory, not one illustrative timing.

1. Why is a.T free but a.T.reshape(-1) not?
2. Given A of shape (1000, 3) and B of shape (2000, 3), write pairwise distances without allocating a (1000, 2000, 3) intermediate.
3. Which of these return views: a[::2], a[a>0], a[[1,2]], a.reshape(2,-1)?
4. counts[idx] += 1 undercounts. Explain and fix.
5. Your 10M × 200 feature matrix takes 16 GB. Halve it with one change.
6. Write attention scores (B, H, Q, K) from Q and K of shape (B, H, T, D) using einsum.
7. Why does np.linalg.solve(A, b) beat np.linalg.inv(A) @ b?

Where to go next

• Pandas — labelled, heterogeneous data on top of these arrays.
• Scikit-learn — models that consume these arrays.
• PyTorch — the same API with autograd and GPUs attached.