Week 4 · Lesson 2

Reduce and expand

Elementwise ops keep a vector a vector. Two more ops change shape: reduce collapses a vector to one number, and expand stretches one number into a vector. Between them they cover sums, means, maxima and normalization, and they're the two halves of almost every layer in a neural network.

reduce is a left fold

reduce %x op=add takes a vector of shape (n,) and returns a scalar. The spec fixes the order exactly:

reduce x op=add  =  ((x[0] + x[1]) + x[2]) + ... + x[n-1]

It starts from x[0], not from 0, and adds the rest left to right, rounding to f32 after every add. op=mul and op=max fold the same way. reduce works on i32 and f32. i32 wraps at every step, as always.

Why the order is in the spec

Real-number addition is associative. f32 addition isn't: every add rounds, and the rounding depends on the running total. The grader compares f32 results bit for bit, so "the sum" has to mean one specific sequence of adds. A left fold from the first element is the simplest one to state, and the easiest to render as a loop.

The same four numbers, folded in two orders. In f32, the spacing between neighboring values near 1e8 is 8, so 1e8 + 1 rounds straight back to 1e8:

Both programs are reduce op=add over {1e8, -1e8, 1, 1}. The first folds to 0 and the second to 2. Each result is the correct left fold of its own program. A compiler that reorders the adds, for example summing pairs in a tree, would get one of them wrong.

Starting from zero

Seeding the sum with 0.0 looks harmless, and it is for almost every input. Except 0.0 + -0.0 is +0.0, so the sum of [-0.0] would come out as +0.0. The spec's fold starts from x[0], so reduce [-0.0] op=add is -0.0. recip tells them apart: -inf versus inf.

A left fold in f32

Write sum32 : float list -> float for a non-empty list: start from the first element, add the rest left to right, and round to f32 after every add with the f32 helper the starter gives you.

add, mul and max

opfolds withedge cases
op=addaddf32 overflow gives inf, and later finite adds keep it inf (a later -inf makes NaN), and i32 wraps
op=mulmul65536 * 65536 * 3 is 0 in i32
op=maxmaxany NaN makes the result NaN, and max(-0.0, 0.0) is +0.0

A one-element vector reduces to that element, with no ops at all. bool vectors can't be reduced.

expand is a broadcast

expand %s shape=[n] takes a scalar and gives a vector of n copies. It's how a scalar meets a vector, since ALU ops never broadcast on their own. Week 4 uses expand on scalars. Week 5 widens it to grow any size-1 axis of a same-rank value. Growing an axis that isn't size 1 is always an error: a (2,) vector can't become [3].

The two ops together give the classic patterns. Here's normalization, y = x / sum(x), written with recip because there's no f32 division op:

%v = load %x                   # (n,)
%s = reduce %v op=add          # ()   the sum
%r = recip %s                  # ()   1 / sum
%e = expand %r shape=[n]       # (n,) 1 / sum, n times
%y = mul %v %e                 # (n,) each x / sum

For x = [1.0 3.0] that's [0.25 0.75]. A mean is a sum times a constant: reduce, then mul by 0.125 for eight elements. Note that %v is used twice, by the reduce and by the final mul. That's fine: the generated code reads x twice whatever you do, because the reduce has to finish before the store loop can start. The real question comes when a computed value (not a plain load) has two uses: recompute it at each use, or store it once. That's the materialization question of lesson 4.4.

Divide by the largest element

Slot 0 holds six f32s. Store y[i] = x[i] * recip(max(x)) into slot 1. Use recip and mul, as written: the tests check that exact rounding, including what NaN and zero maxima do.

The mean of eight

Slot 0 holds eight f32s. Store their sum times 0.125 into slot 1. The starter stores just the sum.

Check yourself

What is reduce %x op=add for %x = stack 100000000.0 -100000000.0 1.0 1.0?

Fold left to right. What's the running total after the first two elements?

1e8 + -1e8 is exactly 0, then 0 + 1 = 1, then 1 + 1 = 2.

And for %x = stack 100000000.0 1.0 1.0 -100000000.0?

In f32 the neighbors of 1e8 are 8 apart. What is 1e8 + 1 after rounding?

1e8 + 1 rounds back to 1e8, twice, and then 1e8 - 1e8 is 0. Same values, different order, different result.

What is reduce %x op=max for %x = stack 1.0 nan 3.0?

reduce folds with week 1's max. What does that do with a NaN operand?

Week 1's max returns NaN if either operand is NaN, so once the NaN enters the fold, every later step is NaN too.

What is reduce %x op=mul for the i32 vector %x = stack 65536 65536 3?

65536 is 216. What's 232 in i32?

65536 * 65536 = 232, which wraps to 0 in i32, and 0 * 3 = 0.

What is reduce %x op=add for %x = stack -0.0?

Where does the fold start, and how many adds does a one-element vector need?

The fold starts from x[0] and there's nothing left to add, so the result is -0.0. A sum seeded with 0.0 would give +0.0.

%v = stack 1 2. What is expand %v shape=[3]?

Which axes is expand allowed to grow?

Week 4 uses expand on scalars, and week 5 lets it grow any size-1 axis. Neither allows a size-2 axis to become size 3, so this is a static error in every week. To repeat a pattern you need the movement ops of week 5.