A pure Python for loop over a million elements takes about a hundred milliseconds; the same computation written with NumPy arrays takes under one millisecond. That two-orders-of-magnitude gap is why almost every machine learning library in Python, from pandas to scikit-learn, is built on NumPy. PyTorch has its own tensor type, but its API follows NumPy's array thinking and it converts to and from ndarray (torch.from_numpy shares memory). This article explains the two concepts that decide whether you can write NumPy efficiently: vectorization (replacing loops with operations on whole arrays) and broadcasting (the rule that lets arrays of different shapes work together). All code was actually run with NumPy 1.23.1 on Python 3.10, on an Apple M4 with 16 GB of RAM; the output is pasted verbatim.
NumPy sits right below pandas. The previous article, Cleaning crawled data with pandas: normalization, deduplication and data types, produced a table with correct types; this one goes down a layer, to where those columns are really numeric arrays. The clean-environment setup from Python environments for data and AI projects: venv, Poetry and version management still applies: run pip install numpy==1.23.1 in a dedicated venv.
Why NumPy is faster than a Python loop
To see why NumPy is fast: a Python list is an array of pointers to float objects scattered across memory, and every addition must check types, look up the method and allocate a result object. A NumPy ndarray is one contiguous block of memory holding numbers of the same type (dtype), plus shape and strides information. NumPy's loops over that block run in compiled C code, without going through the Python interpreter for each element. The paper Array Programming with NumPy (Harris et al., Nature 2020) describes this as the core idea: a few simple array concepts give a compact and powerful way to organize and analyze scientific data, and NumPy became the foundation of the whole scientific Python ecosystem.
Let's measure it with a function that computes the mean squared error (MSE) between true and predicted values over one million elements. The loop version and the vectorized version give the same result:
import time
import numpy as np
rng = np.random.default_rng(42)
# --- A1. Python loop vs vectorization: sum of squared errors ---
n = 1_000_000
y_true = rng.normal(100, 15, n)
y_pred = y_true + rng.normal(0, 5, n)
def mse_loop(a, b):
s = 0.0
for i in range(len(a)):
d = a[i] - b[i]
s += d * d
return s / len(a)
def mse_vec(a, b):
d = a - b
return float(np.mean(d * d))
def best_of(fn, *args, repeat=5):
ts = []
for _ in range(repeat):
t0 = time.perf_counter(); r = fn(*args); ts.append(time.perf_counter() - t0)
return r, min(ts)
r1, t1 = best_of(mse_loop, y_true, y_pred, repeat=3)
r2, t2 = best_of(mse_vec, y_true, y_pred)
print(f"MSE loop : {r1:.6f} ({t1*1000:8.1f} ms)")
print(f"MSE vector : {r2:.6f} ({t2*1000:8.1f} ms)")
print(f"Faster by : {t1/t2:.0f}x")
print("Results match:", np.isclose(r1, r2))MSE loop : 24.962240 ( 117.0 ms)
MSE vector : 24.962240 ( 0.6 ms)
Faster by : 206x
Results match: TrueBoth produce the same MSE, and the vectorized version was 206 times faster in this run (best of 5 for the vector version, best of 3 for the loop version). The exact number changes with the machine and the array size, but the order of magnitude is stable. For arrays of only a few dozen elements, the overhead of calling NumPy functions can outweigh the benefit.
Shape, dtype and memory of a NumPy array
Before writing vectorized NumPy code you need three attributes. shape is the tuple of dimensions; for example (1000, 20) means 1000 samples and 20 features. dtype is the element type, float64 (8 bytes) by default. nbytes is the size of the data block. The common machine learning convention is that the feature matrix has shape (n_samples, n_features), exactly as the scikit-learn documentation describes.
# --- A2. Shape, dtype, memory ---
x = rng.normal(size=(1000, 20))
print("\nshape", x.shape, "dtype", x.dtype, "ndim", x.ndim, "bytes", x.nbytes)
print("float32:", x.astype(np.float32).nbytes, "bytes")shape (1000, 20) dtype float64 ndim 2 bytes 160000
float32: 80000 bytesSwitching to float32 halves the memory, at the cost of lower precision (about 7 decimal digits versus about 16 for float64).
Broadcasting in NumPy: comparing shapes from right to left
When two arrays have different shapes, NumPy does not fail right away; it tries broadcasting. According to the official broadcasting documentation, NumPy compares the shapes of the two arrays element by element, starting from the last (rightmost) dimension and moving left. Two dimensions are compatible when they are equal or when one of them is 1. The arrays do not need the same number of dimensions: a missing dimension is treated as 1 on the left. A dimension of size 1 is "stretched" to match the other without making a real copy.
| Shape A | Shape B | Result | Reason |
|---|---|---|---|
(4, 3) | (3,) | (4, 3) | B is treated as (1, 3) |
(4, 1) | (3,) | (4, 3) | both dimensions are stretched |
(8, 1, 6, 1) | (7, 1, 5) | (8, 7, 6, 5) | compared from the right: 1 and 5, 6 and 1, 1 and 7, missing and 8 |
(4, 3) | (4,) | error | last dimension 3 differs from 4, and neither is 1 |
Let's run the four cases above, plus the fix of adding a new axis. The code blocks labeled B1 to B6 (from here to the end of the View and copy section, except the perf.py benchmark in the middle) are parts of the same script, broadcasting.py: they share the variables X, Z and rng, so paste them in order, starting from the block below (which already contains the imports and setup), and you will get exactly the output pasted here:
import numpy as np
np.set_printoptions(suppress=True, precision=3)
rng = np.random.default_rng(0)
# --- B1. Broadcasting rule: compare shapes from right to left ---
a = np.zeros((4, 3)); b = np.array([1., 2., 3.])
print("(4,3) + (3,) ->", (a + b).shape)
print("(4,1) + (3,) ->", (np.arange(4)[:, None] + np.arange(3)).shape)
print("(8,1,6,1) + (7,1,5) ->", (np.zeros((8, 1, 6, 1)) + np.zeros((7, 1, 5))).shape)
try:
np.zeros((4, 3)) + np.zeros(4)
except ValueError as e:
print("Error:", e)
print("Fix with a column:", (np.zeros((4, 3)) + np.zeros(4)[:, None]).shape)(4,3) + (3,) -> (4, 3)
(4,1) + (3,) -> (4, 3)
(8,1,6,1) + (7,1,5) -> (8, 7, 6, 5)
Error: operands could not be broadcast together with shapes (4,3) (4,)
Fix with a column: (4, 3)The error case is the one you will meet most often. Here (4, 3) + (4,) fails because NumPy aligns on the right: the last dimension of the matrix is 3, and that of the vector is 4. To add one value to each row, turn the vector into a column with [:, None] (equivalent to np.newaxis); its shape becomes (4, 1) and broadcasting stretches it to (4, 3). This technique is worth memorizing: add an axis with None to state explicitly which dimension gets stretched.
Example 1: feature standardization (z-score) with NumPy, no loops
Distance-based and gradient-based algorithms are sensitive to scale: an income of tens of millions would drown out an age of a few dozen. The standard fix is to subtract the mean and divide by the standard deviation, column by column. With X of shape (500, 3), X.mean(axis=0) has shape (3,), and the expression (X - mu) / sd broadcasts (500, 3) with (3,):
# --- B2. Column standardization (z-score): X is (n_samples, n_features) ---
X = np.column_stack([rng.normal(30, 8, 500), # age
rng.normal(15_000_000, 4_000_000, 500), # income (VND)
rng.integers(0, 12, 500)]) # number of orders
mu = X.mean(axis=0); sd = X.std(axis=0) # shape (3,)
Z = (X - mu) / sd # (500,3) - (3,) -> (500,3)
print("\nmean before:", X.mean(axis=0))
print("mean after :", np.abs(Z.mean(axis=0)).round(6), "| std after:", Z.std(axis=0).round(6))mean before: [ 29.785 14723333.38 5.64 ]
mean after : [0. 0. 0.] | std after: [1. 1. 1.]Before standardization the column means range from about 30 to more than 14 million; afterward, the absolute mean is 0 and the standard deviation is 1 in all three columns (printed rounded). In a real project, mu and sd must be computed on the training set and only then applied to the test set, to avoid data leakage.
Example 2: distances between all pairs of points
This problem comes up again and again in clustering, k-NN and similarity search: given P with 5 points and Q with 3 points in 2-D space, compute the 5 × 3 distance matrix. Add axes so that P[:, None, :] has shape (5, 1, 2) and Q[None, :, :] has shape (1, 3, 2); their difference broadcasts to (5, 3, 2). Then square, sum over the last axis and take the square root:
# --- B3. Distances between all pairs of points, no loops ---
P = rng.normal(size=(5, 2)); Q = rng.normal(size=(3, 2))
diff = P[:, None, :] - Q[None, :, :] # (5,1,2)-(1,3,2) -> (5,3,2)
D = np.sqrt((diff ** 2).sum(axis=-1)) # (5,3)
D_loop = np.array([[np.linalg.norm(p - q) for q in Q] for p in P])
print("\ndiff shape:", diff.shape, "| D shape:", D.shape, "| matches loop:", np.allclose(D, D_loop))diff shape: (5, 3, 2) | D shape: (5, 3) | matches loop: TrueThe result matches the version with two nested loops using np.linalg.norm. The price of NumPy here: this expression creates a (5, 3, 2) intermediate array, so memory grows with the product of the number of points and the number of dimensions. The NumPy documentation also warns that temporary arrays can get large. The performance section below measures exactly that.
Example 3: numerically stable softmax per row
Softmax turns a vector of scores into probabilities and appears at the end of most classifiers. Writing exp(z) / sum(exp(z)) directly overflows easily when scores are large. The standard approach is to subtract each row's maximum before taking exp; the mathematical result is unchanged because the numerator and denominator are multiplied by the same factor. The parameter keepdims=True keeps the shape (2, 1) instead of (2,), so the subtraction and division broadcast correctly along rows:
# --- B4. Numerically stable softmax per row ---
logits = np.array([[2.0, 1.0, 0.1], [1000.0, 1001.0, 1002.0]])
def softmax(z):
z = z - z.max(axis=1, keepdims=True) # subtract max: avoid exp overflow
e = np.exp(z)
return e / e.sum(axis=1, keepdims=True)
print("\nsoftmax:\n", softmax(logits).round(4))
with np.errstate(over="ignore", invalid="ignore"):
naive = np.exp(logits) / np.exp(logits).sum(axis=1, keepdims=True)
print("naive softmax (row 2):", naive[1])softmax:
[[0.659 0.242 0.099]
[0.09 0.245 0.665]]
naive softmax (row 2): [nan nan nan]The second row has scores around 1000: the naive version returns nan because exp(1000) exceeds the range of float64 and becomes infinity, while the subtract-max version returns valid probabilities.
Example 4: matrix multiplication for linear regression
With NumPy, a linear regression's predictions for all samples at once are just a matrix product plus a bias: y_hat = X @ w + b. The @ operator calls into an optimized linear algebra library (BLAS), so it is much faster than looping by hand. It is also the basis for computing gradients with a few matrix products in the articles on model training.
# --- B5. Linear regression predictions with matrix multiplication ---
w = np.array([0.5, -0.2, 0.1]); bias = 3.0
y_hat = Z @ w + bias # (500,3)@(3,) -> (500,)
print("\ny_hat shape:", y_hat.shape, "| first 3 values:", y_hat[:3].round(3))y_hat shape: (500,) | first 3 values: [2.882 2.962 3.8 ]Performance: loops, broadcasting and matrix multiplication
Here is a real measurement on the same distance problem, with 2000 points and 500 reference points in 64 dimensions. Three approaches give the same result: two nested loops; broadcasting (2000, 1, 64) - (1, 500, 64); and the identity ||p - q||² = ||p||² + ||q||² - 2·p·q, which reduces the problem to a single matrix product P @ Q.T.
import time
import numpy as np
rng = np.random.default_rng(1)
N, M, D = 2000, 500, 64
P = rng.normal(size=(N, D)); Q = rng.normal(size=(M, D))
def t(fn, repeat=3):
best = 1e9
for _ in range(repeat):
t0 = time.perf_counter(); r = fn(); best = min(best, time.perf_counter() - t0)
return r, best * 1000
def loop():
out = np.empty((N, M))
for i in range(N):
for j in range(M):
out[i, j] = np.sqrt(((P[i] - Q[j]) ** 2).sum())
return out
def bcast():
return np.sqrt(((P[:, None, :] - Q[None, :, :]) ** 2).sum(-1))
def matmul():
# ||p-q||^2 = ||p||^2 + ||q||^2 - 2 p.q
d2 = (P ** 2).sum(1)[:, None] + (Q ** 2).sum(1)[None, :] - 2 * P @ Q.T
return np.sqrt(np.maximum(d2, 0))
a, ta = t(loop, repeat=1)
b, tb = t(bcast)
c, tc = t(matmul)
print(f"Distances {N}x{M} points, {D} dimensions")
print(f"Python loop : {ta:9.1f} ms")
print(f"broadcasting : {tb:9.1f} ms (temp {N*M*D*8/1e6:.0f} MB)")
print(f"matrix multiply : {tc:8.1f} ms")
print("match:", np.allclose(a, b), np.allclose(a, c))Distances 2000x500 points, 64 dimensions
Python loop : 1390.4 ms
broadcasting : 65.3 ms (temp 512 MB)
matrix multiply : 8.0 ms
match: True TrueThe ordering surprises many beginners: broadcasting is about 21 times faster than the loop, but matrix multiplication is about 8 times faster again than broadcasting on this machine, because it does not have to create the (2000, 500, 64) intermediate array. The 512 MB figure is the size of one such intermediate array (2000 × 500 × 64 × 8 bytes); the expression (P[:, None, :] - Q[None, :, :]) ** 2 creates another array of the same size, so real peak memory can be double, around 1 GB.
The 21x and 8x ratios depend on the BLAS library, the NumPy version and the machine (rerunning on another machine can give very different ratios; for example, a faster BLAS makes matrix multiplication pull even further ahead), so rerun it to verify in your own environment. The lesson: choosing a formula with less intermediate memory is the second optimization step after vectorization. The third approach has a price in rounding error: the difference of two nearby numbers can come out very slightly negative, which is why we clamp with np.maximum(d2, 0) before taking the square root.
Views and copies: the trap that corrupts your original data
According to the NumPy documentation on copies and views, basic slicing (X[:, 0], X[:2]) returns a view that shares memory with the original array, while advanced indexing (a list of indices or a boolean array) returns a copy. Modifying a view modifies the original array too:
# --- B6. Views and copies ---
v = X[:, 0]; v[0] = -1; print("\nAfter modifying the view, X[0,0] =", X[0, 0])
c = X[:, 1].copy(); c[0] = -1; print("After modifying the copy, X[0,1] =", round(X[0, 1]))
print("Slice shares memory :", np.shares_memory(X, X[:2, 0]))
print("Fancy indexing shares memory:", np.shares_memory(X, X[[0, 1], 0]))After modifying the view, X[0,0] = -1.0
After modifying the copy, X[0,1] = 20171572
Slice shares memory : True
Fancy indexing shares memory: FalseAfter assigning v[0] = -1 on the view of the first column, X[0, 0] changes with it; modifying the .copy() leaves the original income column unchanged. np.shares_memory gives a definitive answer: slicing shares memory, advanced indexing does not. When you need an independent array, call .copy() explicitly.
When not to vectorize with NumPy, and common mistakes
- Intermediate arrays that are too large: NumPy broadcasting creates temporary arrays sized by the product of the dimensions. With large data, split the work into chunks or switch to a matrix-multiplication formula as above.
- Element-by-element branching logic: if each step depends on the result of the previous one (like a recursive filter), vectorization is hard to apply; consider Numba or rewriting the algorithm.
- Confusing shape
(n,)with(n, 1): in NumPy, a one-dimensional vector is not a column.(n, 1) - (n,)gives an(n, n)matrix instead of(n,), a silent error that is very hard to spot. Alwaysprint(x.shape)when in doubt. - Overflow and loss of precision: a NumPy array of type
int32overflows without raising an error; choose thedtypedeliberately. - Treating
np.vectorizeas a speedup: the NumPy documentation states plainly that this function is only a syntax convenience; it still loops in Python and does not deliver the speed of true vectorization.
NumPy checklist for machine learning
- In NumPy, follow the
(n_samples, n_features)shape convention; printshapeat every step when debugging. - Replace element-by-element loops with whole-array operations, and compare the result with the loop version using
np.allclose. - Use
axisandkeepdims=Trueso that mean, sum and max broadcast back along the right dimension. - Add axes with
[:, None]in NumPy instead of manualreshape. - Estimate the memory of intermediate arrays before broadcasting on large data.
- Choose
float32orfloat64deliberately, and write the reason in the code. - Time with
time.perf_counterortimeit, take the minimum over several runs, and record the machine configuration.
In short, NumPy is powerful not because of short syntax but because it pushes loops down into C code over contiguous memory, while broadcasting lets you express computation between different shapes without writing a loop. Knowing the right-to-left shape rule, adding axes, and watching for intermediate arrays is enough to read most machine learning code written with NumPy.
Next article
With array tools in hand, we return to a more fundamental question: which problems call for machine learning, what the data should look like, and what steps an AI project goes through. Continue with What is machine learning: problems, data and the workflow of an AI project.
