Kiru Lab  /  Foundations: Python as an Instrument  /  Python for Experiments

Arrays and Vectorization

NumPy is the language every later library speaks. Shapes are the grammar.

Code Walkthrough  ·  about 35 minutes

A NumPy array is a contiguous block of numbers plus a shape describing how to interpret it. Operations apply to the whole block at once, in compiled code, without a Python loop. This is not only faster — it is a different way of thinking, where you describe a transformation of a whole dataset rather than of one element.

The same computation, two mindsets
import numpy as np

# Element mindset
out = []
for row in data:                      # data: 1000 rows of 64 numbers
    out.append(sum(r * w for r, w in zip(row, weights)))

# Array mindset
X = np.asarray(data)                  # shape (1000, 64)
w = np.asarray(weights)               # shape (64,)
out = X @ w                           # shape (1000,) — one line, ~100x faster

Shapes are the grammar

Almost every error you will hit in NumPy, PyTorch, or JAX is a shape error, and almost every shape error is legible if you write the shapes down. Adopt the habit now: annotate every array line with its resulting shape in a comment.

(1000, 64) @ (64,) -> (1000,) (1000, 64) @ (64, 10) -> (1000, 10) (32, 8, 128) @ (128, 128) -> (32, 8, 128) # batch, heads, features

The last line is the shape of an attention projection. You already know how to read it.

Broadcasting

When shapes do not match, NumPy will stretch size-1 dimensions to fit rather than erroring. This is enormously convenient and occasionally catastrophic: a (1000, 1) array added to a (1000,) array silently produces a (1000, 1000) array, and your memory disappears.

The broadcast that eats your RAM
a = np.zeros((1000, 1))
b = np.zeros((1000,))
print((a + b).shape)   # (1000, 1000) — one million entries, not one thousand
Structure dictates function

A tensor's shape is its anatomy. Change the arrangement of the same numbers and you change what the operation means — the numbers alone never tell you what a layer does.

Hold on to

  • Vectorized operations describe transformations of whole datasets
  • Annotating shapes turns most bugs into reading comprehension
  • Broadcasting is convenient and silently expensive when misused

Work through

Try each one before opening the solution. Getting it wrong first is most of where the learning happens.

  1. Rewrite a nested-loop matrix multiply as `@` and measure the speedup on 500x500 matrices.
    Hint

    Time both with `time.perf_counter()`; use `np.allclose` to confirm they agree.

    Solution

    The vectorized version is typically 50–500x faster at 500x500, because the work happens in compiled BLAS code rather than in the Python interpreter. Always confirm the results match before believing a speedup — a fast wrong answer is the usual outcome of a rushed vectorization.

    import numpy as np, time

rng = np.random.default_rng(0)
A, B = rng.normal(size=(500, 500)), rng.normal(size=(500, 500))

start = time.perf_counter()
naive = [[sum(A[i, k] * B[k, j] for k in range(500)) for j in range(500)] for i in range(500)]
naive_time = time.perf_counter() - start

start = time.perf_counter()
fast = A @ B
fast_time = time.perf_counter() - start

assert np.allclose(np.array(naive), fast)
print(f"speedup: {naive_time / fast_time:.0f}x")
  2. Deliberately trigger a bad broadcast, then fix it with `reshape` and with `keepdims`. Explain which is clearer.
    Hint

    A (1000, 1) and a (1000,) array broadcast to (1000, 1000).

    Solution

    `reshape(-1)` flattens the column into a 1-D array; `keepdims=True` is what you want when the extra axis came from a reduction you intend to broadcast back against the original. Use `keepdims` when the shape is deliberate and `reshape` when you are fixing an accident — the distinction tells a later reader which it was.

    import numpy as np

a = np.zeros((1000, 1))
b = np.zeros((1000,))
print((a + b).shape)                  # (1000, 1000) — the accident

print((a.reshape(-1) + b).shape)      # (1000,) — flattened

X = np.ones((1000, 5))
means = X.mean(axis=1, keepdims=True) # (1000, 1) on purpose
print((X - means).shape)              # (1000, 5) — intended broadcast
    Check your work

    Paste this after your own code. If it runs without raising, you have it. 

    import numpy as np
assert (np.zeros((1000, 1)) + np.zeros((1000,))).shape == (1000, 1000)
assert (np.zeros((1000, 1)).reshape(-1) + np.zeros((1000,))).shape == (1000,)
print("ok")

In the Bio Mirror

The Cortical Column

Sign in to track your progress through the lab.

← Python for Experiments Seeing Your Data →