Week 5 · Lesson 1

Any rank, flat memory

Until now every value was a scalar or a vector. From this week a value can have any number of axes: a matrix, a stack of matrices, a batch of images. Memory doesn't change at all. Buffers stay one long row of numbers, and the shape is just a way of reading that row.

A shape is a way of counting

A shape is a tuple of positive sizes. () is a scalar, (5,) a vector, (3, 4) a matrix with 3 rows of 4, and (2, 3, 4) two such matrices. The number of axes is the rank.

The elements are laid out in row-major order: the last axis changes fastest. That's the rule from lesson 0.4, and it gives each axis a stride, the distance in elements between neighbors along that axis. For (2, 3, 4) the strides are (12, 4, 1), so element [i][j][k] is number 12i + 4j + k in the flat row. Step through it:

Memory stays rank 1

Buffers, allocs and memory views keep their week 2 shapes: a buffer is (n,), a view is [k] or a single element. So a rank-3 value is always made the same way. Load a vector, then reshape it:

%x  = buffer 24 dtype=i32 slot=0
%v  = load %x                      # (24,)
%t  = reshape %v shape=[2 3 4]     # (2, 3, 4): same numbers, new shape

Storing goes the other way: store needs a value with the view's shape, so you reshape back to [24] first. Because the layout is row-major on both sides, neither reshape moves a single number.

Storing a matrix directly

store %y %m with %m of shape (3, 4) and %y a 12-element buffer is a type error (the oracle reports "the value must match the view's dtype and shape"). Reshape to [12] first. Params also keep their week 1 rule, so param 0 dtype=i32 shape=[2 3] is invalid. Shaped data comes from buffers.

Popping leading axes

Two ops take axes off the front of a shape.

When pop is more than 1, the fold visits the popped axes in row-major order, starting from the first element, exactly like week 4's left fold. For the tensor holding 0..23, reduce pop=2 at column 0 is 0 + 4 + 8 + 12 + 16 + 20 = 60.

Only leading axes

index and reduce work on the front of the shape. To pick or sum along another axis, move it to the front first with permute (lesson 5.2). That's one rule to implement instead of one per axis.

Add the two planes

Slot 0 holds 24 i32s: a (2, 3, 4) tensor in row-major order. Store y[j][k] = x[0][j][k] + x[1][j][k], a (3, 4) result, into slot 1 as 12 numbers. The starter stores plane 0 only.

Unflatten a flat index

Write coord(f, d1, d2, axis): for a row-major tensor of shape [d0 d1 d2], return the coordinate of flat index f along axis 0, 1 or 2. For [2 3 4], flat 23 is [1][2][3].

Check yourself

In a (2, 3, 4) tensor, what is the flat position of element [1][2][3]?

The strides are (12, 4, 1).

12·1 + 4·2 + 3 = 23, the last element. It's the last because every index is at its maximum.

%t has shape (2, 3, 4). What shape is reduce %t op=add pop=1?

Which end of the shape does pop take axes from?

pop=1 folds away the leading axis, so what's left is (3, 4): the two planes added together.

%t is (2, 3, 4) holding 0, 1, ..., 23 in row-major order. What is element 0 of reduce %t op=add pop=2?

The result has shape (4,). Element 0 sums column 0 of all six rows. Column 0 holds every fourth number.

Column 0 is 0, 4, 8, 12, 16, 20, and their sum is 60.

What shape is index %t 1 for %t of shape (2, 3, 4)? Answer like 3 4.

One index drops one leading axis.

It picks plane 1 and drops axis 0, leaving (3, 4).

%m has shape (3, 4) and %y is a 12-element buffer. How do you write %m into %y?

Memory stays rank 1. What does store require of its value's shape?

The value must match the view's shape, and buffers are rank 1. reshape %m shape=[12] moves nothing, because both sides are row-major, and then one store writes all 12.

Is %p = param 0 dtype=i32 shape=[2 3] valid in week 5?

Section 1 of the week 5 spec has one exception to "any rank".

The spec keeps params at rank 0 or 1. Shaped data comes from buffers, loaded and reshaped.