NumPy Vectorization and Broadcasting for Neural Networks

NumPy vectorization expresses many repeated scalar calculations as one array operation. In machine learning, this is more than a speed trick: it gives the implementation the same structure as the underlying linear algebra. Broadcasting complements vectorization by expanding compatible dimensions without making unnecessary copies.

From scalar loops to matrix operations

Suppose a model must calculate a score for 10,000 customers. A beginner may write a loop that processes one customer and then repeats the same instructions 9,999 more times. Vectorization reorganizes those customers into a matrix and describes the repeated work as one matrix operation. The computer still performs arithmetic for every customer, but NumPy sends the work to optimized numerical routines instead of interpreting a Python loop one step at a time.

For vectors \(w,x\in\mathbb{R}^n\), the dot product is \(w^Tx=\sum_jw_jx_j\). With one example per column of \(X\in\mathbb{R}^{n_x\times m}\) and weights shaped (n_x, 1), all scores are \(Z=w^TX+b\), with shape (1, m). Run the following body blocks in order; they reuse the same arrays.

import numpy as np

X = np.array([[1., 2., 3.], [4., 5., 6.]])
w = np.array([[0.2], [-0.1]])
b = 0.3
Z = w.T @ X + b
A = 1 / (1 + np.exp(-Z))
print(Z.shape, A.shape)
print(np.round(Z, 4))
# (1, 3) (1, 3)
# [[0.1 0.2 0.3]]

A complete one-hidden-layer implementation of this is in Building a Shallow Neural Network with NumPy.

The first column gives \(0.2\times1-0.1\times4+0.3=0.1\). The other columns give 0.2 and 0.3, matching the three entries of Z.

Choose one data convention

The body examples store one example per column. Here \(m\) counts examples, \(n_x\) counts input features, and \(n_l\) counts units in layer \(l\). The superscript \([l]\) identifies a layer. An examples-as-rows convention is also valid, provided the equations and operations use it consistently. This choice of axis meaning is separate from row-major or column-major memory storage. The exercises explicitly use examples as rows to practice reading shapes in that layout.

How broadcasting works

Broadcasting can be understood as applying one reusable value along a compatible direction. If a layer has three neurons, its bias has three values—one for each neuron. Those same three values must be added to every example in the batch. A bias shaped (3, 1) makes that intention explicit: preserve the neuron axis and repeat along the example axis.

NumPy compares shapes from the trailing (rightmost) dimensions. Missing leading dimensions are treated as size 1. Two dimensions are compatible when they are equal or one of them is 1. Adding \(b\in\mathbb{R}^{n_l\times1}\) to \(Z\in\mathbb{R}^{n_l\times m}\) therefore reuses each neuron’s bias across the batch, without materializing a repeated bias matrix.

Broadcasting can also hide mistakes. A one-dimensional array with shape (n,) has no explicit row or column orientation. Adding it to a two-dimensional matrix may succeed while producing an unintended result. Prefer (n, 1) or (1, n) and state the intended axis.

Vectorizing loss and gradients

For the sigmoid binary classifier above, use binary labels \(Y\) shaped (1, m). With binary cross-entropy, the derivative with respect to each logit is A - Y; this expression depends on that activation and loss combination. Matrix multiplication aggregates the contributions to dw, and a sum across examples gives db. The loss below is evaluated directly from logits, using the stable form explained in Logistic Regression.

Y = np.array([[0., 1., 1.]])
assert Y.shape == A.shape
cost = np.mean(np.maximum(Z, 0) - Z * Y + np.log1p(np.exp(-np.abs(Z))))
dZ = A - Y
dw = X @ dZ.T / X.shape[1]
db = np.sum(dZ, axis=1, keepdims=True) / X.shape[1]
print("cost", round(float(cost), 6))
print("dw", dw.shape, np.round(dw.ravel(), 6))
print("db", db.shape, np.round(db.ravel(), 6))
# cost 0.632297
# dw (2, 1) [-0.550675 -0.901419]
# db (1, 1) [-0.116915]

Axis, reduction, and keepdims

A reduction combines values along the axis named by axis. With examples in columns, sum(axis=1) adds across examples for each feature. By default that axis is removed; keepdims=True retains it with length 1. For this two-dimensional array, the result is a column vector:

row_sums = X.sum(axis=1)
kept_sums = X.sum(axis=1, keepdims=True)
print(row_sums.shape, row_sums)
print(kept_sums.shape, kept_sums)
normalized = X / kept_sums
print(np.round(normalized, 4))
print(normalized.sum(axis=1))
# (2,) [ 6. 15.]
# (2, 1) [[ 6.]
#  [15.]]
# [[0.1667 0.3333 0.5   ]
#  [0.2667 0.3333 0.4   ]]
# [1. 1.]

The denominator has shape (2, 1), so every entry in the first row is divided by 6 and every entry in the second by 15. Each resulting row sums to 1. This arithmetic example uses positive row sums; division by a zero sum would need separate handling.

Shape bugs to watch for

• Using X.T in one function but not in the next.
• Allowing labels to become shape (m,) instead of (1,m).
• Using elementwise * where matrix multiplication @ is required.
• Calling reshape with the correct size but the wrong semantic order.
• Reducing over the feature axis instead of the example axis.
• Depending on accidental broadcasting from a one-dimensional array.

A reliable debugging workflow

The following check uses the same two-feature, three-example batch. It checks shapes explicitly before comparing a loop calculation with the vectorized result. Shape checks matter because np.allclose can itself broadcast unequal shapes. The finite-value check can catch numerical failures, but neither check proves that an implementation is correct for every input.

assert X.ndim == 2
assert w.shape == (X.shape[0], 1)
assert Y.shape == Z.shape == (1, X.shape[1])
assert np.all(np.isfinite(A))
Z_loop = np.zeros_like(Z)
for k in range(X.shape[1]):
    Z_loop[0, k] = b
    for f in range(X.shape[0]):
        Z_loop[0, k] += w[f, 0] * X[f, k]
assert Z_loop.shape == Z.shape
print(np.allclose(Z_loop, Z, rtol=1e-12, atol=1e-12))
# True

The L-layer forward and backward passes are written out in Deep Neural Networks: Architecture and Backpropagation.

The comparison agrees within the stated tolerance on this batch. Exercise 2 repeats the comparison on a larger matrix before measuring runtime. Timing needs repeated measurements because system activity, allocation, and library warm-up can affect a single call.

Vectorization is easiest when the mathematical convention, array layout, and code all agree. Once those contracts are explicit, broadcasting becomes a precise tool instead of a source of mysterious results.

The same training stages are implemented with PyTorch tensors, modules, and autograd in Building a PyTorch Training Loop with nn.Module, Dataset, and DataLoader.

Exercises

1. Find the shape bug. In this exercise, examples are rows. The subtraction runs without error, but should produce one residual per example. Identify the bug and state the shape of diff.

You should get: a square matrix where a column vector was intended.

import numpy as np
y_hat = np.random.rand(100, 1)      # predictions, (100, 1)
y     = np.random.rand(100)        # labels loaded as 1-D, (100,)
diff  = y_hat - y
print(diff.shape)                  # (100, 100)

2. Vectorize a loop. Here X has 500 example rows and 64 feature columns, and W maps 64 features to 10 outputs. Rewrite the loop as one matrix expression, check agreement, and compare runtimes.

Check agreement within a stated tolerance, then report the measured runtime ratio without assuming a minimum speedup.

import numpy as np
from timeit import repeat
rng = np.random.default_rng(0)
X = rng.standard_normal((500, 64))
W = rng.standard_normal((64, 10))

def loop(X, W):
    out = np.zeros((X.shape[0], W.shape[1]))
    for i in range(X.shape[0]):
        for j in range(W.shape[1]):
            out[i, j] = np.dot(X[i], W[:, j])
    return out

# Warm up both implementations while checking their results.
a = loop(X, W)
b = X @ W
assert a.shape == b.shape
print(np.allclose(a, b, rtol=1e-10, atol=1e-12))
# True
t_loop = min(repeat(lambda: loop(X, W), repeat=5, number=20)) / 20
t_array = min(repeat(lambda: X @ W, repeat=5, number=20)) / 20
print(f"loop {t_loop:.6g}s, array {t_array:.6g}s, ratio {t_loop/t_array:.1f}x")

3. keepdims, why. Here rows are examples and columns are classes. Softmax exponentiates the scores and divides each row by its sum, producing class probabilities for that example. Compute it for a \((3,5)\) matrix twice: once with keepdims=True in both the max and the sum, once without. Compare the shapes and the row sums.

You should get: one version whose rows sum to 1, and one that either errors or produces something that does not.

import numpy as np
Z = np.arange(15, dtype=float).reshape(3, 5)
good = np.exp(Z - Z.max(axis=1, keepdims=True))
good /= good.sum(axis=1, keepdims=True)
print(good.shape, np.round(good.sum(axis=1), 6))
try:
    bad = np.exp(Z - Z.max(axis=1))
except ValueError:
    print("ValueError: (3, 5) and (3,) are incompatible")

square = np.arange(1, 10, dtype=float).reshape(3, 3)
wrong = np.exp(square - square.max(axis=1))
wrong /= wrong.sum(axis=1)
print("row sums", np.round(wrong.sum(axis=1), 4))
print("column sums", np.round(wrong.sum(axis=0), 4))
# (3, 5) [1. 1. 1.]
# ValueError: (3, 5) and (3,) are incompatible
# row sums [  0.8727  17.5285 352.0696]
# column sums [3.679747e+02 2.479400e+00 1.670000e-02]