Week 5 study guide

Week 5 lifts the rank limit. Values can have any rank while memory stays flat, and four movement ops (reshape, permute, pad, shrink) rearrange them without computing anything. A matmul or a convolution becomes a few movement ops and one reduce. This page recaps the week, collects the rules worth memorizing, and walks the lab one step at a time.

The week in one page

5.1 Any rank, flat memory

Shapes can have any number of axes, stored in row-major order. Buffers, allocs and memory views stay rank 0 or 1: load a rank-1 value, then reshape it, and reshape back to rank 1 before a store. index pops leading axes, and reduce pop=m folds the m leading axes away.

5.2 Reshape and permute

reshape reads the same row-major sequence in a new shape. permute reorders the axes: shape[k] = x.shape[order[k]]. A reshape is not a transpose, even when the shapes match.

5.3 Pad, shrink, expand, stack

pad adds zeros on each side of each axis. A value shrink takes a window with one offset per axis. The offsets may be computed, but must be provably in range. expand grows size-1 axes, and stack joins values of one shape along a new leading axis.

5.4 Matmul and conv from movement

The syllabus GEMM reshapes A to [K M 1] and B to [K 1 N], broadcasts both to [K M N], multiplies, and reduces the leading axis. A 1D convolution is a pad, a stack of shifted shrinks, a multiply by the broadcast kernel, and a reduce.

5.5 Rangeify movement ops

"Your element at index i" now takes one index per axis. Reshape is a flatten and an unflatten with mul, div and mod. Permute reorders the indices, shrink adds offsets, expand reads size-1 axes at 0, and pad clamps the index into range and zeroes the value with where.

5.6 Lab

Same commands as week 4, with every rank limit lifted except on memory and params. The suite has 94 tests, and weeks 1 to 4 must stay green.

Cheat sheet

op	result shape	element at `i`
`reshape x shape=[...]`	that shape (same element count)	the same position in row-major order
`permute x order=[p0 ...]`	`shape[k] = x.shape[p_k]`	`x` at the index whose axis `p_k` is `i_k`
`pad x pads=[l0 r0 ...]`	each axis grows by `l + r`	`x[i - l]` inside, zero outside
`shrink x o0 ... shape=[...]`	that shape	`x[i + o]`, with offsets provably in range
`expand x shape=[...]`	that shape	`x` with each size-1 axis read at 0
`stack a0 ... a(n-1)`	`(n,) + shape`	`a_{i0}[i1, ...]`
`index x j0 ...`	drops one leading axis per index	`x[j0, ..., i0, ...]`
`reduce x op= pop=m`	drops m leading axes	left fold over the leading axes, in row-major order

formula	for shape `(d0, d1, d2)`
flat index (row-major)	`f = (i0 · d1 + i1) · d2 + i2`
unflatten	`i2 = f mod d2`, `i1 = (f div d2) mod d1`, `i0 = f div (d1 · d2)`
matmul by movement	A `[M K]` → `[M K 1]` → permute `[1 0 2]` → expand `[K M N]`. B `[K N]` → `[K 1 N]` → expand. Then mul, then `reduce pop=1` → `[M N]`

Mistakes people make

Each of these is a bug that was planted in the reference compiler to prove the suite catches it. The week 5 test guide lists every test with the bugs known to fail it.

Reshaping in column-major order. Row-major means the last axis changes fastest. Caught by broadcast_column_sums_back, conv2d_same and expand_rank3_middle.
Applying the inverse of a permutation. shape[k] = x.shape[order[k]], not the other way round. Caught by permute_rank3_all_orders, permute_rank4 and permute_inverse_roundtrip.
Swapping a pad's left and right, or filling with something other than zero. Caught by pad_2d_mixed, pad_3d and pad_bool_false.
Getting a shrink offset's sign wrong. The window starts at +o. Caught by pad_shrink_roundtrip_3d and conv1d_same.
An expand that reads index i instead of 0 on a size-1 axis. Caught by broadcast_column_sums_back and conv1d_same.
Reducing the trailing axes, or folding in column-major order. pop=m folds the leading m axes, in row-major order. Caught by reduce_rowmajor_order and broadcast_column_sums_back.
Treating a same-rank reshape as a no-op. [2 3] to [3 2] keeps the flat order, which is not the transpose. Caught by r_reshape_same_rank_then_transpose.
Off-by-one pad masks in render. The valid range is l ≤ i < l + d. Caught by pad_2d and r_conv2d_16x16.
Wrong strides in flatten or unflatten. Caught by gemm_movement_3x4x5 and r_column_normalize.
Skipping the movement-op static checks. Element counts, permutations, pad lengths, shrink bounds and stack shapes are all checked before anything runs. Caught by err_reshape_count, err_permute_duplicate, err_pad_length, err_shrink_past_end and err_stack_mixed_shape.

Stuck on the lab?

Build in this order, and don't move on until the check passes. Save each check as a file and run it through your compiler and the oracle (python _reference/uirc.py run FILE or render FILE).

Shapes of any rank. Done when index and reshape agree on row-major order:
```
%v = stack 0 1 2 3 4 5 6 7 8 9 10 11
%m = reshape %v shape=[3 4]
%x = index %m 2 1
%t = assert %x 9
```
If you get 5, your reshape is column-major (2 + 1 · 3).
Type the movement ops. Done when bad shapes exit 2. Check: %v = stack 1 2, %e = expand %v shape=[3] (a size-2 axis can't grow), and the suite's err_permute_duplicate, err_pad_length and err_shrink_unprovable.

Evaluate them in run. Done when permute, pad and shrink agree with the spec:

%v  = stack 1 2 3 4 5 6
%m  = reshape %v shape=[2 3]
%tr = permute %m order=[1 0]
%a  = reshape %tr shape=[6]
%w  = stack 1 4 2 5 3 6
%t  = assert %a %w
%s  = shrink %m 0 1 shape=[2 2]
%f  = reshape %s shape=[4]
%w2 = stack 2 3 5 6
%u  = assert %f %w2

Then every week 5 run test should pass.

Index tuples in rangeify. Done when a reshape round trip ([12] to [3 4] and back, then stored) renders as a single loop whose index math works out to the loop counter. The oracle prints ((r / 4 % 3) * 4 + r % 4) % 12, which a simplifier can prove equals r. Your text doesn't have to match, only the results.
Each op as index arithmetic. Done when the suite's pad_2d and r_transpose_64x64 pass bit for bit. For pad, clamp the index into range first, then mask the value with where: reading past a buffer is undefined behavior even if the mask discards the value.
Reduce over several axes. Done when reduce_rowmajor_order and r_reduce_pop2_64x64 pass: one accumulator seeded with the first element, one loop over a flat counter covering the other T - 1 combinations (T is the product of the popped axis sizes), unflattened in row-major order.
Optional: clean up the index math. Run your week 1 simplifier on the lowered indices. Constant folding and bounds-based mod elimination remove most of what flatten and unflatten leave behind.

Reshape is not a transpose

A test with constant values, or a symmetric square matrix, can't tell a reshape from a transpose. Test with distinct values everywhere, like 0 1 2 ... 11.

Check yourself

With %v = stack 0 1 2 ... 11 and %m = reshape %v shape=[3 4], what is index %m 2 1?

Row-major: row 2 starts at element 2 · 4.

Row 2 starts at 8, and column 1 is one further: 9.

What shape does permute with order=[2 0 1] give for a value of shape [2 3 4]? (Write the three sizes.)

shape[k] = x.shape[order[k]].

shape[0] = x.shape[2] = 4, shape[1] = x.shape[0] = 2, shape[2] = x.shape[1] = 3: [4 2 3].

In that permuted value, element (3, 1, 2) is the original element at which flat position of the [2 3 4] value?

Axis 0 of the result is axis 2 of the original, axis 1 is axis 0, axis 2 is axis 1. So the original index is (1, 2, 3).

The original index is (1, 2, 3), and its flat position is 1 · 12 + 2 · 4 + 3 = 23, the last element.

How many elements does pad %v pads=[2 1] give for a 3-element vector?

Two zeros on the left, one on the right.

3 + 2 + 1 = 6: [0 0 x0 x1 x2 0].

Which expand is valid?

Only size-1 axes grow, and a non-scalar keeps its rank.

[2 1] to [2 3] grows the size-1 axis. A size-2 axis can't become 3, and a rank-1 value can't become rank 2 (reshape it to [2 1] first).

With %m the [2 2] matrix [[1 2] [3 4]], what is reduce %m op=add (pop defaults to 1)?

pop=1 folds away the leading axis: for each column, sum down the rows.

Column 0 is 1 + 3 = 4 and column 1 is 2 + 4 = 6, so [4 6].

Unflatten the flat index 17 in a [3 6] shape. (Write row and column.)

Row is 17 div 6, column is 17 mod 6.

17 div 6 = 2 and 17 mod 6 = 5, so (2, 5): 2 · 6 + 5 = 17.

More practice: the week 5 quiz review.

Go deeper

tinygrad and its documentation. The movement ops here are tinygrad's, and reading how it tracks views shows how a real compiler keeps reshapes and permutes free.
Row- and column-major order. The two conventions, and the stride formulas behind flatten and unflatten.
NumPy's ndarray. Shapes, strides and views in the library most ML code started from.
PyTorch broadcasting semantics. The implicit rules that expand makes explicit.
CS231n: Convolutional Neural Networks. The im2col view of convolution as a matrix multiply, which lesson 5.4 builds from movement ops.
Anatomy of High-Performance Matrix Multiplication (Goto and van de Geijn), and BLIS. How a fast GEMM is organized: a preview of weeks 6 to 8.
Hypothesis. Property tests that generate random shapes and orders and check that permute and reshape round trip.
QCheck. The same kind of property test in OCaml: random shapes and orders, shrunk to the smallest failure.
Python and OCaml, deeper. This week's idioms in your language: strides, unraveling, and index expressions as UOps.

Challenges

Flip without a flip op. The syllabus calls FLIP a composite op. Build a reversed 4-element vector using only movement ops and arithmetic you have (hint: a matmul with an anti-diagonal matrix built by stack works). Check the result with the oracle.
A 2D transpose by reshape alone? Prove that no sequence of reshapes can transpose a [2 3] matrix with distinct values, then confirm with the oracle that one permute does.
Simplify the index math. Take the oracle's rendered C for a reshape round trip and work out by hand, with interval reasoning from lesson 1.4, which div and mod operations are redundant. Check your claim by changing the program's size and seeing whether the same simplification still holds.