NumPy Broadcasting & Vectorization
Writing high performance vectorized array operations using NumPy shape broadcasting rules.
The Power of Vectorization
Python for loops are slow because the Python interpreter checks object data types and performs dynamic dispatch on every iteration.
Vectorized NumPy Operations execute compiled C linear algebra instructions across contiguous memory blocks in parallel:
# SLOW: Python For Loop (~1.2 seconds for 1,000,000 elements)
result = [a + b for a, b in zip(list_a, list_b)]
# FAST: NumPy Vectorization (~0.005 seconds for 1,000,000 elements!)
result = array_a + array_b
Vectorization speeds up mathematical array operations by 100 to 200 times!
The 2 Rules of NumPy Broadcasting
Broadcasting allows arithmetic operations on arrays of different shapes without making unnecessary data memory copies.
When operating on two arrays $A$ and $B$, NumPy compares their shapes trailing dimension by trailing dimension (from right to left):
Rule 1: Dimensions are equal (dim_A == dim_B).
Rule 2: One of the dimensions is 1 (dim_A == 1 or dim_B == 1).
If neither rule holds, NumPy raises a ValueError: operands could not be broadcast together.
EXAMPLE 1: Matrix + Row Vector
Array A (2D): 4 x 3
Array B (1D): 3
---------------------
Broadcast: 4 x 3 (Compatible! Array B is stretched across 4 rows!)
EXAMPLE 2: Outer Operation using np.newaxis
Array A (1D): 4 x 1 (using A[:, np.newaxis])
Array B (1D): 1 x 3 (using B[np.newaxis, :])
---------------------
Broadcast: 4 x 3 (Outer Product Matrix!)
Useful Broadcasting Patterns in ML
1. Subtracting Feature Means (Standardization)
# X shape: [1000, 50] (1000 samples, 50 features)
# mean shape: [50]
mean = np.mean(X, axis=0)
X_centered = X - mean # Stretches mean across 1000 rows automatically!
2. Pairwise Euclidean Distance Matrix
# A shape: [100, 10], B shape: [50, 10]
# Expand dims: A -> [100, 1, 10], B -> [1, 50, 10]
diff = A[:, np.newaxis, :] - B[np.newaxis, :, :] # Shape: [100, 50, 10]
dists = np.sqrt(np.sum(diff**2, axis=2)) # Shape: [100, 50]
Say this out loud
NumPy Broadcasting and Vectorization execute array math at C compiled speeds. Broadcasting compares shapes from right to left, stretching dimensions equal to 1 to match compatible target shapes without allocating extra memory. Vectorization replaces Python loops with SIMD hardware instructions.
Followups to expect
- Does broadcasting copy memory array data? No, broadcasting operates virtually using stride tricks (
stride = 0for broadcast dimensions), performing math without copying memory. - What is
np.einsum? Einstein summation convention notation for specifying tensor contractions, transpositions, and matrix products concisely (np.einsum('bik,bkj->bij', A, B)).
Check yourself
What fundamental condition determines whether two NumPy array dimensions are compatible for broadcasting?