Coding for ML

NumPy Broadcasting & Vectorization

Writing high performance vectorized array operations using NumPy shape broadcasting rules.

🟢 beginner5 min readcodingnumpy
NumPy Broadcasting and Vectorization enables fast array operations without explicit C or Python loops. Broadcasting defines rules for operating on arrays of different shapes by implicitly stretching smaller dimensions. Vectorized operations execute low level BLAS linear algebra instructions, achieving 100x speedup over Python loops.

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

  1. Does broadcasting copy memory array data? No, broadcasting operates virtually using stride tricks (stride = 0 for broadcast dimensions), performing math without copying memory.
  2. 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

Question 1 of 3

What fundamental condition determines whether two NumPy array dimensions are compatible for broadcasting?

More in Coding for ML

See all →
Implement Linear Regression from Scratch5 minImplement Self-Attention from Scratch5 minPandas: GroupBy, Merge, Window5 min