Delete the loop: five vectorization patterns

The first lesson of this module said why a Python loop over an array is slow, and it is worth restating because it is the whole motivation. Every iteration is an interpreter step: follow a pointer to an object, ask that object what addition means, allocate a result, store it. The numbers are sitting in one contiguous block and nothing is walking that block -- the interpreter is walking a list of instructions instead.

Deleting the loop is not a micro-optimization. It is writing down what the operation is, once, so that a compiled routine can walk the block. The expression is usually shorter than the loop too, and that turns out to be the smaller benefit.

But you cannot delete a loop you have not read. Four questions, in this order, and between them they name the pattern:

• What does the index touch? If i appears only as a[i] and out[i] -- the same position on both sides -- then nothing crosses between elements. Pattern one.
• Is there an accumulator? A variable set before the loop and updated inside it is a reduction, and the loop it lives in is the axis. Pattern two.
• Is there an if? A condition asked one element at a time is a mask over all of them. Pattern three.
• Is the loop nested? If the two indices reach into different arrays, one operand repeats and that is broadcasting, pattern four. If both indices reach into both arrays and the innermost loop accumulates, the shared index is summed away, and that is a matrix multiply, pattern five.

Most real loops are two or three of those stacked. Read from the inside out and name each one as you meet it.

The simplest loop reads element i, does some arithmetic, and writes element i. Nothing crosses between elements, so nothing needs an index:

for i in range(len(a)):          a * 2 + 1
    out[i] = a[i] * 2 + 1

Every arithmetic operator on an array is a ufunc -- a universal function, meaning one operation applied elementwise to everything, in compiled code. +, -, *, / and ** are ufuncs, and so are np.exp, np.sqrt, np.abs and np.maximum.

Two things follow, and both are worth being explicit about.

The rank does not matter. A ufunc applies to every element however many axes there are, so a doubly nested loop over a matrix collapses to the same expression a single loop would have needed:

x.shape == (2, 3)

for i in range(2):                x * 3 + 2
    for j in range(3):
        out[i][j] = x[i][j] * 3 + 2

A scalar joins in. The 2 in a * 2 has shape (), which matches everything, so mixing arrays and plain numbers needs no ceremony at all.

What it costs is worth knowing. Every operation in a chain builds a whole new array: a * 2 + 1 allocates the doubled array, then allocates the shifted one and throws the first away. Almost always a fine trade, and where it is not, the in-place forms -- a *= 2 -- write back into the buffer you already have.

Set a variable before the loop, update it inside, read it afterwards. That is an accumulator, and every accumulator is a reduction:

total = 0                        a.sum()
for value in a:
    total += value

sum, mean, max and min cover nearly all of them. A running best tracked with an if is still a reduction -- max replaces sum and nothing else about the loop matters.

The only real difficulty is the axis, and there is one rule: the axis you name is the axis that disappears. The loop tells you which one that is. Look at what the accumulator is indexed by; the other index is the one being eaten:

x.shape == (4, 3)

for j in ...:                    one total per column
    for i in ...:                -> x.sum(axis=0)
        out[j] += x[i][j]

for i in ...:                    one total per row
    for j in ...:                -> x.sum(axis=1)
        out[i] += x[i][j]

Both loops touch every element exactly once. They differ only in which index the accumulator carries -- and that index is the one that survives, so the other one is the one you name.

keepdims is for the case where the reduction feeds an elementwise operation back onto the array it came from. Without it the reduced axis is gone, and what is left gets padded on the left against the wrong axis:

x.shape == (3, 4)

x.mean(axis=1)                   -> (3,)     a flat row
x.mean(axis=1, keepdims=True)    -> (3, 1)   a column

x - x.mean(axis=1)                  raises
x - x.mean(axis=1, keepdims=True)   centers every row

That first line raises here only because 3 and 4 differ. On a square array it raises nothing and centers the wrong direction, which is the version that costs people an afternoon.

a > 3 is not a question about a. It is an array of bools with the same shape as a:

a       == [1, 5, 3, 8, 3]
a > 3   == [False, True, False, True, False]

There are three things to do with a mask, and reaching for the wrong one is the commonest mistake in this pattern:

• Count. (a > 3).sum() -- summing bools counts the Trues, because True adds one.
• Filter. a[a > 3] keeps the values that passed, as a shorter array. The shape is gone, which is either what you wanted or a bug.
• Choose. np.where(cond, x, y) keeps the shape and picks a value per position.

Which one you need is decided by what the if was doing:

count = 0                        (a > 3).sum()
for v in a:
    if v > 3: count += 1

total = 0                        a[a > 3].sum()
for v in a:
    if v > 3: total += v

out = []                         np.where(a < 0, 0, a)
for v in a:
    if v < 0: out.append(0)
    else:     out.append(v)

The argument order of np.where is the order the if is written in: the condition, then the value when it holds, then the value when it does not. Getting those two the wrong way round is the single most productive typo in array code, because it runs, it has the right shape, and it keeps exactly the values you meant to throw away.

Some conditions have a shorter name. np.maximum(a, 0) is a floor, np.minimum(a, 10) is a ceiling, np.clip(a, 0, 10) is both. They are the same thing with less to get backwards.

What it costs: np.where evaluates both branches for every element and then selects. An if in a loop evaluates one. That is nearly always the right trade, and it stops being one when a branch is invalid rather than merely unwanted -- dividing by the thing the condition was there to keep you away from still divides.

Two nested loops, and one of the arrays inside is indexed by only one of them. That array is the one that repeats, and broadcasting is how you say so:

for i in range(len(x)):          x + b
    for j in range(len(x[0])):
        out[i][j] = x[i][j] + b[j]

b is indexed by j, the columns, which are the last axis. A shorter shape is padded on the left, so a bare b already lines up and there is nothing to write.

for i in range(len(x)):          x * s[:, None]
    for j in range(len(x[0])):
        out[i][j] = x[i][j] * s[i]

s is indexed by i, the rows, which are the first axis. Padding would put it against the last one, so it has to be made a column first. That is the entire decision in this pattern: which index does the small operand know about, and is that the last axis? If it is, do nothing; if it is not, add the length-one axes that push it into place.

The other shape of nested loop is the one where the two indices reach into two different arrays. Then neither operand repeats along its own index, both repeat along the other one, and what comes out is a grid:

for i in range(len(a)):          a[:, None] - b[None, :]
    for j in range(len(b)):
        out[i][j] = a[i] - b[j]

A column against a row. Which operand gets the column decides which index runs down the rows, and swapping them transposes the answer -- which for a difference also flips every sign. Distance matrices, attention scores and Gram matrices are all built this way.

What it costs: the grid is real memory even though the operands are not. Two arrays of n numbers produce n by n results, and for a few thousand points that is already larger than the data that made it.

This is the pattern the other four have been building towards, and it has a signature you can spot without reading the arithmetic: an index that appears in both operands and is summed over.

total = 0                        a @ b
for k in range(len(a)):
    total += a[k] * b[k]

for i in ...:                    W @ v
    total = 0
    for k in ...:
        total += W[i][k] * v[k]
    out[i] = total

for i in ...:                    A @ B
    for j in ...:
        total = 0
        for k in ...:
            total += A[i][k] * B[k][j]
        out[i][j] = total

The indices split into two kinds. i and j pick out a position in the answer, so they survive as its axes. k is shared by both arrays and the sum runs over it, so it disappears. Written as shapes, that is the whole rule:

(m, k) @ (k, n)  ->  (m, n)      the k on the inside is consumed

Which is why the order is not negotiable. A @ B and B @ A are different operations, and on square operands both of them run.

* is not @. The star is the elementwise ufunc from pattern one and has no idea what a shared dimension is. Python gained the at sign as a separate operator precisely so that the two could stop being confused, and it is worth reading every * in matrix code as "no sum happens here".

This is also where the whole lesson pays off. A matrix multiply is not a loop that happens to be compiled -- it is a call into a linear algebra routine written against the machine's cache hierarchy and vector instructions, and on a GPU into a kernel written against its tensor cores. Every dense layer in every network you will build is X @ W + b, so the arithmetic that goes fastest is the arithmetic that runs most.

Five patterns cover most of the loops in numerical code, and it would be a bad lesson that left you believing they cover all of them.

Sequential dependence. When step i needs the result of step i - 1, there is no single expression, because the operation is defined in terms of its own output:

out[0] = a[0]
for i in range(1, len(a)):
    out[i] = 0.9 * out[i - 1] + a[i]

Some special cases have library forms -- np.cumsum and np.cumprod are running totals and running products written in compiled code. A genuine recurrence has none, and neither does a training loop over epochs or an agent stepping an environment. Those loops are correct as loops. The work inside one iteration is what gets vectorized.

The grid you cannot afford. The pairwise construction from pattern four turns n points into n by n entries. Past some size the loop is merely slow and the vectorized version does not run at all. The fix is neither one nor the other: take the points a block at a time, so the grid is a chunk by n and the loop runs a handful of times instead of n times.

The line nobody can read. Four patterns stacked into one expression with three Nones and a keepdims is not obviously better than six clear lines, especially in code that runs once at startup. Vectorize the inner loop that runs a million times, and leave the outer one alone if it makes the file readable.

And whichever you choose, check it. Keep the loop. Write the expression next to it. Run both on a small array and compare. Then, and only then, delete the loop.

assert np.allclose(slow(x), fast(x))

A vectorized expression fails silently far more readily than a loop does, because broadcasting will happily give you an answer of the wrong shape without complaining. The loop you are deleting is the best test you will ever have of the expression replacing it, and it is free right up until the moment you delete it.