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.
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.
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
| op | folds with | edge cases |
|---|---|---|
op=add | add | f32 overflow gives inf, and later finite adds keep it inf (a later -inf makes NaN), and i32 wraps |
op=mul | mul | 65536 * 65536 * 3 is 0 in i32 |
op=max | max | any 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.