Week 5 · Lesson 3
Pad, shrink, expand, stack
Four more ways to rearrange without computing: grow a tensor with zeros, cut a window out of it, repeat it along a size-1 axis, or join several into one. With these, lesson 5.2's two, and index from lesson 5.1, you can write down most of the data movement in a neural network.
pad adds zeros
pad %x pads=[l0 r0 l1 r1 ...] takes two numbers per axis: how many zeros go before it and how many after. A 2D value has four, [top bottom left right]. The new cells hold the dtype's zero (0, 0.0 or false), and element [i][j] of the result is x[i - l0][j - l1] when that's inside x.
pad [[1 2] [3 4]] pads=[0 1 2 0] adds a row below and two columns on the left: shape (3, 4), rows 0 0 1 2, 0 0 3 4, 0 0 0 0. The pads must be at least 0; to cut, use shrink.
A ring of zeros
Slot 0 holds a 2x2 f32 grid. Store the 4x4 grid with it in the middle and zeros all around into slot 1. The starter's grid is the right size with the zeros in the wrong place.
shrink cuts a window
shrink %x o0 o1 shape=[s0 s1] takes one offset per axis and a window size: element [i][j] is x[i + o0][j + o1]. It's week 2's memory view, applied to a value. The first operand decides which one you get: memory gives a view, a value gives a window.
The offsets can be any i32 expression, but as with views they must be provably in range, using the same interval check as week 2's views (lesson 2.3): for every axis, lo ≥ 0 and hi + s ≤ size. A window of 2 over 4 elements at offset mod (max %i 0) 3 is fine, since the offset is in [0, 2] and 2 + 2 ≤ 4. At mod (max %i 0) 4 the offset can reach 3, and 3 + 2 > 4, so it's rejected. A raw loaded %i could be anything, so it's rejected too.
Padding by one on every side and then shrinking at offsets 1 1 with the original shape gives the original back. Shrinking anywhere else sees some of the zeros. That pair, pad and a sliding window, is how a padded convolution is built (lesson 5.4).
expand repeats along size-1 axes
Week 4's expand took a scalar. Now it also takes a value of the same rank, and grows only its size-1 axes: (3, 1) can become (3, 4), and element [i][j] reads x[i][0]. A size-2 axis can't become 3, and the rank can't change. To broadcast a vector across rows, give it a size-1 axis first with reshape: (4,) becomes (1, 4), then expands to (3, 4).
Add a bias to every row
Slot 0 holds a 3x4 f32 matrix and slot 1 a bias of 4 values. Store y[i][j] = x[i][j] + b[j] into slot 2. The starter adds b[0] everywhere.
stack joins along a new axis
stack %a %b needs operands of one shape and puts a new axis in front: two (2, 2) grids become a (2, 2, 2) tensor, and index on the new axis gets each one back. In week 1 it only stacked scalars, but now it stacks anything, as long as the shapes match.
Check yourself
What shape is pad %x pads=[1 0 0 2] for a (2, 2) %x?
The pairs are per axis: (1, 0) for rows, (0, 2) for columns.
Rows grow by 1 + 0 to 3, columns by 0 + 2 to 4.
%x = [[0 1] [2 3]]. What is element [2][1] of pad %x pads=[1 0 0 2]?
Subtract the left pads: [2 - 1][1 - 0].
It reads x[1][1], which is 3. Element [0][0] would be a padding zero, and so would any column past 1.
%c has shape (3, 1) and holds 0, 1, 2. What is element [2][3] of expand %c shape=[3 4]?
The size-1 axis is read at 0, whatever the output index is.
It reads c[2][0], which is 2. Every column of row 2 is 2.
%a = stack 1 2. Is expand %a shape=[3] valid?
Which axes can expand grow?
An axis of size 2 can't become 3. The oracle says: "can't expand [2] to [3]: same rank, and only size-1 axes grow".
%v has shape (4,) and %i is a loaded i32. Which shrink %v %o shape=[2] is valid?
The rule is lo ≥ 0 and hi + 2 ≤ 4. What's each offset's interval?
mod (max %i 0) 3 is in [0, 2], and 2 + 2 ≤ 4. With mod ... 4 it can reach 3, and a raw load can be anything, so both are rejected. Offsets can be symbolic as long as they're provable.
What shape is stack %x %x for a (2, 2) %x? Answer like 2 3.
stack puts one new axis in front, sized by the number of operands.
Two operands of shape (2, 2) give (2, 2, 2).