%uML Compiler from Scratch
Part III · Week 5 · 68 min read

9. A Tensor library and your first MNIST

You can now compile any shape graph into C, cut it into kernels, and differentiate it. But nobody wants to write shape graphs. Typing A.reshape((M, 1, K)).expand((M, N, K)) every time you mean A @ B is how compiler people punish themselves, not how machine-learning people work. This chapter builds the thin layer that turns the compiler into something that feels like PyTorch: a Tensor class with broadcasting, reductions over any axes, matrix multiplication, convolution, softmax and cross-entropy; a handful of layers; two optimizers. Then we use it to train a classifier on MNIST, the course's week-5 milestone, and look closely at the 32 C kernels that make up one training step. None of this needs a new compiler feature. Every operation in this chapter is a few lines of graph building, and that is the point: if your IR is right, the frontend is small.

What you'll build

  • uopc/tensor.py: a lazy Tensor that wraps a shaped UOp graph, with NumPy-style broadcasting, movement ops, reductions over arbitrary axes, matmul, conv2d, softmax, log_softmax, argmax, one_hot and cross_entropy.
  • uopc/nn.py: Linear, Conv2d, LayerNorm, Embedding, a parameters() walker, and the SGD and AdamW optimizers.
  • examples/mnist_mlp.py: a 784-128-10 MLP that trains to 94% test accuracy in about nine seconds.

Where this fits. Chapter 7 taught the compiler to split a graph into kernels and Chapter 8 taught it to differentiate one. This chapter is the user-facing API on top of both. Chapter 10 then makes the kernels it produces fast.

9.1 Why a lazy Tensor

A NumPy array is eager: a + b computes and stores a new array immediately. An eager design is simple, but it hides the program from the compiler. If y = (x * 2 + 1).relu().sum(1) ran eagerly, it would be four separate loops with three temporary arrays, and no later pass could fuse them, because by the time sum runs the temporaries already exist.

Our Tensor is lazy. It is a thin wrapper around a UOp, and each operator builds a bigger UOp graph. Nothing runs until somebody asks for numbers (.numpy(), .item(), or .realize()). At that moment the whole pending graph goes through the pipeline from Chapter 7: kernelize, rangeify, render, compile (or fetch from cache), run.

Laziness buys three things, each of which this chapter relies on:

  1. Fusion. The compiler sees the whole expression and decides where to cut it into kernels.
  2. Autodiff on graphs. gradient() from Chapter 8 needs the forward graph as data. An eager library would have to record a tape; a lazy one already is the tape.
  3. Whole-step JIT caching. A training step builds the same graph shape every iteration, so its compiled plan can be reused (_structure_key in kernelize.py).

The price is that errors and costs show up late. A shape mismatch is caught when you build the graph (the UOp type checker runs at construction), but a slow kernel only shows up at realize().

Here is a session, run with DEBUG=1 so that every realize prints one line:

>>> x = Tensor(np.arange(6, dtype=np.float32).reshape(2, 3))
>>> y = (x * 2 + 1).relu().sum(1)
>>> y
<Tensor (2,) f32>
>>> y.uop.op, y.uop.src[0].op, y.uop.src[0].arg
Ops.REDUCE Ops.WHERE None
>>> y.numpy()
realize: 1 kernels, schedule 6.1 ms, run 0.1 ms
[ 9. 27.]
>>> y.uop.op, y.uop.arg
Ops.BUFFER ParamArg(slot=2, dtype=dtypes.f32, shape=(2,), addrspace=None)

Before .numpy(), y is a REDUCE over a WHERE (that is relu) over arithmetic over a BUFFER. Building it ran no C at all. After .numpy() the tensor has been rebound: y.uop is now a BUFFER that holds the result, so asking again costs nothing, and any graph built on y later starts from data rather than recomputing the whole chain.

9.2 The Tensor class

The class starts with a constructor that accepts either a UOp (the internal path, used by every operator) or anything NumPy can turn into an array (the user path):

uopc/tensor.py · Tensor.__init__
  def __init__(self, data, dtype: DType | None = None, requires_grad: bool = False):
    if isinstance(data, UOp): self.uop = data
    else:
      arr = np.asarray(data)
      if dtype is None: dtype = dtypes.f32 if arr.dtype.kind == "f" else dtypes.bool if arr.dtype.kind == "b" else dtypes.i32
      arr = arr.astype(np_dtype(dtype))
      buf = new_buffer(dtype, arr.size, arr)
      self.uop = buf.reshape(arr.shape) if arr.shape != buf.shape else buf
    self.requires_grad = requires_grad
    self.grad: Tensor | None = None

User data is copied into a fresh BUFFER UOp (new_buffer in kernelize.py attaches the NumPy array to the UOp itself, so the memory lives exactly as long as the node does; see Chapter 7). BUFFERs are flat, so a 2-D array becomes BUFFER(6).reshape((2, 3)). dtypes are inferred the obvious way: floats become f32, bools bool, everything else i32.

Realization is where the frontend hands over to the compiler:

uopc/tensor.py · Tensor.realize_all
  @staticmethod
  def realize_all(*ts: Tensor):
    todo = [t for t in ts if t.uop.op is not Ops.BUFFER and not (t.uop.op is Ops.RESHAPE and base_buffer(t.uop).op is Ops.BUFFER and prod(t.shape) == base_buffer(t.uop).arg.shape[0])]
    if not todo: return
    bufs = _realize([t.uop for t in todo], Tensor.fuse_reduce, Tensor.opt)
    for t, b in zip(todo, bufs):
      t._graph = t.uop   # keep the graph: loss.item() before loss.backward() must still have gradients
      t.uop = b.reshape(t.shape) if b.shape != t.shape else b

Three details matter. First, tensors that are already data (a BUFFER, or a reshape of a whole BUFFER) are skipped; realizing them would launch a useless copy kernel. Second, realize_all realizes many tensors in one call. That is how an optimizer step updates every parameter, every moment estimate and the loss in a single schedule: one structural cache lookup, one batch of kernels. Realizing them one by one would work, but it would compute shared subexpressions (the gradients) once per tensor. Third, before rebinding a tensor to its result BUFFER, realize_all keeps the graph it came from in t._graph, so that loss.item() followed by loss.backward(...) still has a graph to differentiate (the Pitfalls tell how the first version got this wrong).

assign is the only "mutation" in the library, and it is not really one:

uopc/tensor.py · Tensor.assign
  def assign(self, x: Tensor) -> Tensor:
    """Make this tensor (e.g. a parameter) hold x's value. Functional: rebinds to a new buffer."""
    assert x.shape == self.shape and x.dtype == self.dtype, f"assign {x.shape} to {self.shape}"
    self.uop = x.uop
    return self

It rebinds the Python object to a new graph. Buffers are never written in place by user code; a parameter update builds p - lr * g and then makes p point at that expression. When the expression is realized, p.uop becomes the new BUFFER and the old one is garbage-collected along with its UOp. This functional style keeps autodiff sound (no node ever changes its value) and lets the kernelizer reason purely about graphs.

9.3 Broadcasting

Elementwise ops in the IR are strict: ADD demands that both operands have the same shape. The type checker from Chapter 1 says so, loudly. Users expect NumPy's broadcasting instead: (2, 3) + (3,) adds the vector to every row. The rules are:

  1. Line up the shapes from the right; pad the shorter one with leading 1s.
  2. In each position the sizes must be equal, or one of them must be 1.
  3. The result takes the larger size in each position.
uopc/tensor.py · _broadcast_shape
def _broadcast_shape(*shapes):
  n = max(len(s) for s in shapes)
  shapes = [(1,) * (n - len(s)) + tuple(s) for s in shapes]
  out = []
  for dims in zip(*shapes):
    d = max(dims)
    if any(x not in (1, d) for x in dims): raise ValueError(f"can't broadcast {shapes}")
    out.append(d)
  return tuple(out)

Turning a shape into a broadcast shape takes exactly two primitive movement ops: a RESHAPE to add the leading 1s, and an EXPAND to stretch the 1s:

uopc/tensor.py · Tensor.broadcast_to
  def broadcast_to(self, shape) -> Tensor:
    shape = tuple(shape)
    if self.shape == shape: return self
    u = self.uop
    if len(u.shape) < len(shape): u = u.reshape((1,) * (len(shape) - len(u.shape)) + u.shape)
    return Tensor(u.expand(shape) if u.shape != shape else u)

and every binary operator funnels through one helper:

uopc/tensor.py · Tensor._binop
  def _binop(self, op: Ops, x, reverse=False) -> Tensor:
    # numpy-style promotion for python scalars: int tensor (op) 0.5 is a float computation
    if isinstance(x, float) and not self.dtype.is_float: return self.float()._binop(op, x, reverse)
    x = self._const(x)
    a, b = (x, self) if reverse else (self, x)
    shape = _broadcast_shape(a.shape, b.shape)
    return Tensor(UOp(op, (a.broadcast_to(shape).uop, b.broadcast_to(shape).uop)))

So a + b with a of shape (2, 3) and b of shape (3,) becomes ADD(a, EXPAND(RESHAPE(b, (1, 3)), (2, 3))). In the REPL:

>>> c = a + b
>>> c.uop.op, [s.op for s in c.uop.src]
Ops.ADD [Ops.RESHAPE, Ops.EXPAND]
>>> c.uop.src[1].op, c.uop.src[1].shape, c.uop.src[1].src[0].op, c.uop.src[1].src[0].shape
Ops.EXPAND (2, 3) Ops.RESHAPE (1, 3)

(a itself shows up as a RESHAPE because it is a flat BUFFER viewed as (2, 3).) Does the EXPAND cost anything? Not a byte. Rangeify (Chapter 6) turns EXPAND into "use index 0 on the stretched axis", so the generated C simply never looks at ridx0 when it reads b:

void e_2_3(float* restrict data0, float* restrict data1, float* restrict data2) {
  for (int ridx0 = 0; ridx0 < 2; ridx0++) {
    for (int ridx1 = 0; ridx1 < 3; ridx1++) {
      int alu1 = ((ridx0*3)+ridx1);
      float val2 = data0[alu1];
      float val3 = data1[ridx1];
      data2[alu1] = (val2+val3);
    }
  }
}

Key idea: broadcasting is free because it is indexing, not copying. A broadcast operand is read through an index expression that ignores the broadcast loop. The same fact is what makes matmul, one-hot and argmax cheap to express below.

The rest of the elementwise vocabulary

With _binop in hand, the arithmetic operators are one-liners. Several are derived rather than primitive, because the IR deliberately has a small ALU (ADD, MUL, MAX, CMPLT, CMPNE, RECIP, EXP2, LOG2, SQRT, SIN, WHERE, CAST):

uopc/tensor.py:129-153
  def __add__(self, x): return self._binop(Ops.ADD, x)
  def __radd__(self, x): return self._binop(Ops.ADD, x, True)
  def __mul__(self, x): return self._binop(Ops.MUL, x)
  def __rmul__(self, x): return self._binop(Ops.MUL, x, True)
  def __neg__(self): return self * -1
  def __sub__(self, x): return self + -self._const(x)
  def __rsub__(self, x): return self._const(x) + -self
  def __truediv__(self, x): return self * self._const(x).reciprocal()
  def __rtruediv__(self, x): return self._const(x) * self.reciprocal()
  def __lt__(self, x): return self._binop(Ops.CMPLT, x)
  def __gt__(self, x): return self._binop(Ops.CMPLT, x, True)
  def __ge__(self, x): return (self < x).logical_not()
  def __le__(self, x): return (self > x).logical_not()
  def eq(self, x): return self._binop(Ops.CMPNE, x).logical_not()
  def ne(self, x): return self._binop(Ops.CMPNE, x)
  def logical_not(self): return self._binop(Ops.CMPNE, True)
  def maximum(self, x): return self._binop(Ops.MAX, x)
  def minimum(self, x): return -((-self).maximum(-self._const(x)))
  def where(self, a, b) -> Tensor:
    a = a if isinstance(a, Tensor) else Tensor.full((), a)
    b = b if isinstance(b, Tensor) else Tensor.full((), b, a.dtype)
    shape = _broadcast_shape(self.shape, a.shape, b.shape)
    return Tensor(UOp(Ops.WHERE, (self.broadcast_to(shape).uop, a.broadcast_to(shape).uop, b.broadcast_to(shape).uop)))
  def cast(self, d: DType) -> Tensor: return Tensor(self.uop.cast(d))
  def float(self) -> Tensor: return self.cast(dtypes.f32)

The first line of _binop (shown above) is NumPy-style type promotion: an integer tensor combined with a Python float is cast to f32 first, so Tensor([1, 2, 3]) * 0.5 is [0.5, 1.0, 1.5] (see the Pitfalls for what happened without it). Subtraction is addition of a negation, division is multiplication by a reciprocal, > is < with the operands swapped, == is "not !=". Why not add SUB and DIV to the IR? Because every op you add must be handled by the renderer, the symbolic simplifier, the bounds analysis and autodiff. Deriving a - b as a + b * -1 costs nothing at runtime (gcc sees x + y*-1.0f and emits a subtraction) and saves four rules in four files.

Transcendentals follow the same philosophy:

uopc/tensor.py:155-170
  def reciprocal(self): return Tensor(self.uop.recip())
  def exp2(self): return Tensor(UOp(Ops.EXP2, (self.uop,)))
  def log2(self): return Tensor(UOp(Ops.LOG2, (self.uop,)))
  def exp(self): return (self * LOG2E).exp2()
  def log(self): return self.log2() * LN2
  def sqrt(self): return Tensor(UOp(Ops.SQRT, (self.uop,)))
  def rsqrt(self): return self.sqrt().reciprocal()
  def sin(self): return Tensor(UOp(Ops.SIN, (self.uop,)))
  def cos(self): return (self + math.pi / 2).sin()
  def square(self): return self * self
  def abs(self): return self.maximum(-self)
  def relu(self): return (self > 0).where(self, 0.0)
  def sigmoid(self): return ((-self).exp() + 1).reciprocal()
  def tanh(self): return (self * 2).sigmoid() * 2 - 1
  def gelu(self): return 0.5 * self * (1 + (math.sqrt(2 / math.pi) * (self + 0.044715 * self * self * self)).tanh())
  def silu(self): return self * self.sigmoid()

The only exponential in the IR is EXP2, because that is what hardware has (exp2f on CPUs, ex2.approx on NVIDIA GPUs). exp(x) is exp2(x * log2(e)) and log(x) is log2(x) * ln 2. You will see the constant 1.4426950408889634f (that is log2 e) all over the generated C in this chapter. relu is a WHERE, sigmoid is built from exp and a reciprocal, and tanh from sigmoid. Every one of these gets a gradient for free, because autodiff only needs rules for the primitives.

9.4 Movement ops

The movement methods are thin conveniences over the four primitive movement ops from Chapter 6. reshape accepts a -1 and works out the missing size, permute accepts negative axes, transpose swaps two axes, unsqueeze inserts a 1, flatten merges trailing axes. Indexing with ints and step-1 slices is a SHRINK followed by a RESHAPE that drops the integer-indexed axes:

uopc/tensor.py · Tensor.__getitem__
  def __getitem__(self, idx) -> Tensor:
    """Ints and step-1 slices only."""
    idx = idx if isinstance(idx, tuple) else (idx,)
    arg, squeeze = [], []
    for d, i in enumerate(idx):
      n = self.shape[d]
      if isinstance(i, int):
        i = i % n
        arg.append((i, i + 1)); squeeze.append(d)
      else:
        b, e, step = i.indices(n)
        assert step == 1, "only step-1 slices"
        arg.append((b, max(b, e)))
    arg += [(0, s) for s in self.shape[len(arg):]]
    t = self.shrink(arg)
    return t.reshape(tuple(s for d, s in enumerate(t.shape) if d not in squeeze))

pad and flip call the composite constructions from composite.py (built from STACK, RESHAPE and SHRINK; see Chapter 6). Concatenation reuses them: each piece is zero-padded into its place along the axis, and the padded pieces are added.

uopc/tensor.py · Tensor.cat
  @staticmethod
  def cat(ts: list[Tensor], dim: int = 0) -> Tensor:
    """Concatenate via pad + add: each piece is padded into place, then summed."""
    dim = dim % ts[0].ndim
    total = sum(t.shape[dim] for t in ts)
    out, off = None, 0
    for t in ts:
      p = [(0, 0)] * t.ndim
      p[dim] = (off, total - off - t.shape[dim])
      piece = t.pad(p)
      out = piece if out is None else out + piece
      off += t.shape[dim]
    return out

Adding zeros sounds wasteful, but after rangeify each padded piece becomes a bounds-checked WHERE, so cat costs a few comparisons per element, not extra memory traffic.

9.5 Reductions over arbitrary axes

The IR's REDUCE has a deliberately narrow contract: it reduces the last n axes. Users want x.sum((0, 2)). The bridge is a permutation: move the kept axes to the front, the reduced axes to the back, then pop:

uopc/tensor.py · Tensor._reduce
  def _reduce(self, op: Ops, axis=None, keepdim=False) -> Tensor:
    axes = list(range(self.ndim)) if axis is None else [a % self.ndim for a in (axis if isinstance(axis, (tuple, list)) else [axis])]
    keep = [i for i in range(self.ndim) if i not in axes]
    t = self.permute(keep + axes)
    r = Tensor(t.uop.reduce(op, len(axes)))
    if keepdim: r = r.reshape(tuple(1 if i in axes else s for i, s in enumerate(self.shape)))
    return r

For x of shape (2, 3, 4) and axis=(0, 2), keep = [1], so we permute by (1, 0, 2) to shape (3, 2, 4) and REDUCE the last 2 axes, giving shape (3,). With keepdim=True a reshape puts the 1s back where the reduced axes were, which is what makes x - x.max(-1, keepdim=True) broadcast correctly.

>>> x = Tensor(np.arange(24, dtype=np.float32).reshape(2, 3, 4))
>>> x.sum((0, 2)).numpy(), x.sum((0, 2)).uop.src[0].arg
[ 60.  92. 124.] (1, 0, 2)

The permutation, like the expand, is free. It only changes which loop variable multiplies which stride:

void r_3(float* restrict data0, float* restrict data1) {
  for (int ridx0 = 0; ridx0 < 3; ridx0++) {
    float acc0;
    acc0 = 0.0f;
    for (int ridx1 = 0; ridx1 < 2; ridx1++) {
      for (int ridx2 = 0; ridx2 < 4; ridx2++) {
        float val1 = acc0;
        float val2 = data0[(((ridx1*12)+(ridx0*4))+ridx2)];
        acc0 = (val1+val2);
      }
    }
    float val3 = acc0;
    data1[ridx0] = val3;
  }
}

The outer loop is the kept axis (stride 4 in memory), the two inner loops are the reduced axes (strides 12 and 1). The accumulator acc0 lives inside the output loop and outside the reduce loops, exactly the placement rule from Chapter 6.

The derived reductions are short:

uopc/tensor.py · Tensor.mean
  def mean(self, axis=None, keepdim=False):
    s = (self if self.dtype.is_float else self.float()).sum(axis, keepdim)
    return s * (s.numel() / self.numel())
uopc/tensor.py · Tensor.var
  def var(self, axis=None, keepdim=False):
    return (self - self.mean(axis, keepdim=True)).square().mean(axis, keepdim)

mean sums in f32 (casting integer tensors first) and multiplies by numel(result) / numel(input), a Python float that becomes a CONST. var is the mean squared deviation. min is a negated max of the negation (the IR has only ADD and MAX reductions). softmax, log_softmax and logsumexp get their own section below.

9.6 Matrix multiplication is broadcast, multiply, reduce

Here is the most important three-line function in the library:

uopc/tensor.py · Tensor.matmul
  def matmul(self, w: Tensor) -> Tensor:
    """(..., M, K) @ (..., K, N): broadcast to (..., M, N, K), multiply, reduce K."""
    x = self if self.ndim > 1 else self.reshape(1, -1)
    wt = w.reshape(-1, 1) if w.ndim == 1 else w
    wt = wt.transpose()                                     # (..., N, K)
    x = x.reshape(x.shape[:-1] + (1, x.shape[-1]))          # (..., M, 1, K)
    wt = wt.reshape(wt.shape[:-2] + (1,) + wt.shape[-2:])   # (..., 1, N, K)
    out = (x * wt).sum(-1)
    if self.ndim == 1: out = out.reshape(out.shape[:-2] + out.shape[-1:])
    if w.ndim == 1: out = out.reshape(out.shape[:-1])
    return out

Read it as geometry. A matrix product C[i, j] = sum_k A[i, k] * B[k, j] has three indices, so lay the computation out on an (M, N, K) box. A has no j index, so it is reshaped to (M, 1, K) and broadcast along N. B has no i index; transposed to (N, K) and reshaped to (1, N, K) it broadcasts along M. Multiplying gives the full box of products A[i, k] * B[k, j], and summing the last axis (K) collapses the box onto the (M, N) face, which is C.

Matrix multiplication as broadcast, multiply, reduce A of shape (M, 1, K) is expanded along N and B transposed of shape (1, N, K) is expanded along M; together they fill an (M, N, K) box of products. Summing the box along K collapses it onto the (M, N) face, which is C. One output element C[i, j] is the line through the box along K that pairs row i of A with row j of B transposed. expand N expand M sum over K A (M, 1, K) row i Bᵀ (1, N, K) row j N M K products (M, N, K) C[i, j] C (M, N) the box is a loop nest, never a buffer

The box is never stored. It is a description of the iteration space, and rangeify turns it into a loop nest in which each multiply happens once, right before it is accumulated. For a (4, 8) @ (8, 3) product with the optimizer switched off (OPT=0), the graph and the code are:

[('REDUCE', (4, 3), (Ops.ADD, 1)), ('MUL', (4, 3, 8), None), ('EXPAND', (4, 3, 8), (4, 3, 8)),
 ('RESHAPE', (4, 1, 8), (4, 1, 8)), ('RESHAPE', (4, 8), (4, 8))]

void r_4_3(float* restrict data0, float* restrict data1, float* restrict data2) {
  for (int ridx0 = 0; ridx0 < 4; ridx0++) {
    for (int ridx1 = 0; ridx1 < 3; ridx1++) {
      float acc0;
      acc0 = 0.0f;
      for (int ridx2 = 0; ridx2 < 8; ridx2++) {
        float val1 = acc0;
        float val2 = data0[((ridx0*8)+ridx2)];
        float val3 = data1[((ridx2*3)+ridx1)];
        acc0 = (val1+(val2*val3));
      }
      float val4 = acc0;
      data2[((ridx0*3)+ridx1)] = val4;
    }
  }
}

That is the textbook triple loop, produced by a library function that never mentions a loop. (The first line above follows src[0] down from the result; the max error against NumPy was 2.4e-07.) Look at the two loads in the inner loop. A is read at ridx0*8 + ridx2: consecutive k are adjacent in memory. B is read at ridx2*3 + ridx1: consecutive k are a whole row apart. For a 3-wide matrix that does not matter; for a 1024-wide one it is the difference between 1.9 and 8 GFLOP/s, and it is the first thing Chapter 10 fixes.

Batched matmul falls out of the same code. All leading axes (...) ride along untouched, so (B, M, K) @ (B, K, N) broadcasts to (B, M, N, K) and reduces K. The test test_ops[matmul_batched] checks a @ b.transpose() on (2, 3, 4) tensors against NumPy. linear is x @ w + b, with the bias broadcast across the batch.

9.7 Convolution as im2col of shrinks

A 2-D convolution with a KH x KW kernel computes, for every output pixel, a dot product between the kernel and the KH x KW window of input under it, summed over input channels. The classic way to make convolution a matrix multiply is im2col: gather every window into a row of a big "patches" matrix, then multiply by the flattened weights. Libraries usually implement im2col as a copy. We get it from movement ops alone:

uopc/tensor.py · Tensor.conv2d
  def conv2d(self, w: Tensor, b: Tensor | None = None, padding: int = 0, stride: int = 1) -> Tensor:
    """x (N, C, H, W), w (O, C, KH, KW). im2col from shrinks, then a matmul."""
    x = self.pad(((0, 0), (0, 0), (padding, padding), (padding, padding))) if padding else self
    N, C, H, W = x.shape
    O, _, KH, KW = w.shape
    OH, OW = (H - KH) // stride + 1, (W - KW) // stride + 1
    cols = []
    for ky in range(KH):
      for kx in range(KW):
        win = x.shrink(((0, N), (0, C), (ky, ky + stride * (OH - 1) + 1), (kx, kx + stride * (OW - 1) + 1)))
        if stride > 1: win = win._stride2d(stride, OH, OW)
        cols.append(win)                                     # (N, C, OH, OW) each
    patches = cols[0].stack_with(*cols[1:]) if len(cols) > 1 else cols[0].unsqueeze(0)   # (KH*KW, N, C, OH, OW)
    patches = patches.permute(1, 3, 4, 2, 0).reshape(N, OH, OW, C * KH * KW)
    out = patches @ w.reshape(O, C * KH * KW).T              # (N, OH, OW, O)
    out = out.permute(0, 3, 1, 2)
    return out + b.reshape(1, O, 1, 1) if b is not None else out

For each kernel offset (ky, kx) there is a SHRINK of the (padded) input that selects the pixels that offset touches for every output position; each has shape (N, C, OH, OW). STACKing the KH*KW shrinks gives (KH*KW, N, C, OH, OW). Permuting to (N, OH, OW, C, KH*KW) and merging the last two axes gives the patches matrix (N, OH, OW, C*KH*KW), whose column order (channel-major, then ky, then kx) matches w.reshape(O, C*KH*KW). One matmul later, a permute puts channels back in NCHW order.

Strided convolution needs every s-th window. With no strided SHRINK in the IR, _stride2d pads the window up to a multiple of s, reshapes (OH*s) into (OH, s), and takes index 0 of the s axis:

uopc/tensor.py · Tensor._stride2d
  def _stride2d(self, s: int, OH: int, OW: int) -> Tensor:
    N, C, H, W = self.shape
    t = self.pad(((0, 0), (0, 0), (0, OH * s - H), (0, OW * s - W)))
    return t.reshape(N, C, OH, s, OW, s)[:, :, :, 0, :, 0]

Max pooling uses the same reshape trick without the pad: (H, W) becomes (H/k, k, W/k, k) and the two k axes are max-reduced.

Is it right? A (1, 2, 5, 5) input, (3, 2, 3, 3) weights, padding=1, checked against a NumPy loop:

(1, 3, 5, 5) max err 1.4305115e-06

Is it fast? Here is the reduce kernel, with OPT=0. It is worth reading closely, because it is the plain output of a compiler that has not yet learned anything about convolutions:

void r_1_5_5_3(float* restrict data0, float* restrict data1, float* restrict data2) {
  for (int ridx0 = 0; ridx0 < 5; ridx0++) {
    int alu1 = (ridx0*5);
    int alu2 = ((ridx0+5)%6);
    int alu3 = (alu2*5);
    int alu4 = ((ridx0+7)%6);
    int alu5 = (alu4*5);
    for (int ridx1 = 0; ridx1 < 5; ridx1++) {
      int alu6 = (ridx1+5);
      int alu7 = (alu6%6);
      int alu8 = (alu7%5);
      int alu9 = ((ridx1+7)%6);
      int alu10 = (alu9%5);
      int alu11 = (ridx1+17);
      int alu12 = (alu11%6);
      int alu13 = ((ridx1+19)%6);
      int alu14 = (ridx1+29);
      int alu15 = (alu14%6);
      int alu16 = (alu15%5);
      int alu17 = ((ridx1+31)%6);
      int alu18 = (alu17%5);
      for (int ridx2 = 0; ridx2 < 3; ridx2++) {
        float acc0;
        acc0 = 0.0f;
        for (int ridx3 = 0; ridx3 < 18; ridx3++) {
          float val19 = acc0;
          int alu20 = (ridx3%9);
          int alu21 = (ridx3/9);
          int alu22 = (alu21*25);
          int alu23 = (alu22+((alu2%5)*5));
          float val24 = data0[(alu23+alu8)];
          float val25 = data0[(alu23+ridx1)];
          float val26 = data0[(alu23+alu10)];
          int alu27 = (alu22+alu1);
          float val28 = data0[(alu27+(alu12%5))];
          float val29 = data0[(alu27+ridx1)];
          float val30 = data0[(alu27+(alu13%5))];
          int alu31 = (alu22+((alu4%5)*5));
          float val32 = data0[(alu31+alu16)];
          float val33 = data0[(alu31+ridx1)];
          float val34 = data0[(alu31+alu18)];
          float val35 = data1[((((ridx2*18)+(alu21*9))+(((ridx3/3)%3)*3))+(ridx3%3))];
          acc0 = (val19+(((alu20<1)?(((alu6/6)<1)?0.0f:(((alu7/5)<1)?(((((alu1+alu8)+25)/30)<1)?0.0f:((((alu3+alu8)/25)<1)?val24:0.0f)):0.0f)):((alu20<2)?(((((alu1+ridx1)+25)/30)<1)?0.0f:((((alu3+ridx1)/25)<1)?val25:0.0f)):((alu20<3)?(((alu9/5)<1)?(((((alu1+alu10)+25)/30)<1)?0.0f:((((alu3+alu10)/25)<1)?val26:0.0f)):0.0f):((alu20<4)?((((alu11/6)%2)<1)?0.0f:(((alu12/5)<1)?val28:0.0f)):((alu20<5)?val29:((alu20<6)?(((alu13/5)<1)?val30:0.0f):((alu20<7)?((((alu14/6)%2)<1)?0.0f:(((alu15/5)<1)?((((alu5+alu16)/25)<1)?val32:0.0f):0.0f)):((alu20<8)?((((alu5+ridx1)/25)<1)?val33:0.0f):(((alu17/5)<1)?((((alu5+alu18)/25)<1)?val34:0.0f):0.0f)))))))))*val35));
        }
        float val36 = acc0;
        data2[(((ridx0*15)+(ridx1*3))+ridx2)] = val36;
      }
    }
  }
}

Three things are visible. The 18-long reduce loop is C * KH * KW flattened, and ridx3 % 9 picks which of the nine kernel taps this iteration belongs to: that is the STACK, lowered to a chain of WHERE (alu20<1 ? ... : alu20<2 ? ...). The %6, /30 and /25 terms are the composite PAD, which Chapter 6 builds from concatenation with zeros; each tap carries its own bounds check. And all nine candidate loads are hoisted above the select, so every iteration performs nine loads to use one. The kernel is correct and general, and roughly nine times more memory-hungry than it should be. A convolution-aware compiler would make the tap index a real loop (or unroll it) instead of selecting from a STACK; that is exercise 9.10.

9.8 Softmax, log-softmax and cross-entropy

Softmax turns a vector of scores into a probability distribution: softmax(x)_i = exp(x_i) / sum_j exp(x_j). Computed literally in f32, it overflows as soon as a score exceeds about 88 (exp(89) is larger than the largest float). The cure is an identity: for any constant m, exp(x_i - m) / sum_j exp(x_j - m) equals exp(x_i) / sum_j exp(x_j), because the exp(-m) factors cancel. Choosing m = max_j x_j makes every exponent at most 0, so nothing overflows, and the largest term is exactly 1, so the sum is at least 1 and nothing underflows to a 0/0.

Classification losses want the log of the probability of the right class. Taking the log of a softmax is numerically poor (a tiny probability rounds to 0 and its log to -inf), so we compute it directly:

log_softmax(x)_i = (x_i - m) - log(sum_j exp(x_j - m)).

uopc/tensor.py · Tensor.softmax
  def softmax(self, axis=-1):
    e = (self - self.max(axis, keepdim=True)).exp()
    return e / e.sum(axis, keepdim=True)
uopc/tensor.py · Tensor.log_softmax
  def log_softmax(self, axis=-1):
    m = self - self.max(axis, keepdim=True)
    return m - m.exp().sum(axis, keepdim=True).log()
uopc/tensor.py · Tensor.logsumexp
  def logsumexp(self, axis=-1):
    m = self.max(axis, keepdim=True)
    return (self - m).exp().sum(axis, keepdim=True).log() + m

Cross-entropy for a batch of B examples with integer labels y_b is L = -(1/B) * sum_b log_softmax(z_b)[y_b]. Picking "the entry at index y_b" is a gather, and gathers with data-dependent indices are exactly what the IR cannot differentiate (_rule for INDEX refuses). So we turn the label into a one-hot row and use a multiply-and-sum, which is an ordinary reduce:

uopc/tensor.py · Tensor.one_hot
  def one_hot(self, n: int) -> Tensor:
    return self.unsqueeze(-1).eq(Tensor.arange(n)).float()
uopc/tensor.py · Tensor.cross_entropy
  def cross_entropy(self, labels: Tensor) -> Tensor:
    """self: (B, C) logits, labels: (B,) i32."""
    return -(self.log_softmax(-1) * labels.one_hot(self.shape[-1])).sum(-1).mean()

one_hot is itself a broadcast: labels of shape (B,) become (B, 1), arange(n) has shape (n,), and comparing them broadcasts to (B, n) with a 1 exactly where the column equals the label:

>>> Tensor(np.array([2, 0, 3], np.int32)).one_hot(4).numpy()
[[0. 0. 1. 0.]
 [1. 0. 0. 0.]
 [0. 0. 0. 1.]]
>>> logits.cross_entropy(labels).item(), numpy_reference
0.31853973865509033 0.3185397696491856

How many kernels does softmax become? REDUCE is a materialization point (Chapter 7), so: one kernel for the row max, one for the row sum of exponentials (which reads the max), and one elementwise kernel for the output. The test suite pins this down:

tests/test_w05_tensor.py · test_kernelize_fuses_elementwise_and_splits_reductions
def test_kernelize_fuses_elementwise_and_splits_reductions():
  a = Tensor(r(64, 64))
  plan, _ = make_plan([((a * 2 + 1).relu() * 3).uop])
  assert len(plan.kernels) == 1                         # a chain of elementwise ops: one kernel
  plan, _ = make_plan([a.softmax(-1).uop])
  assert len(plan.kernels) == 3                         # max, sum, and the elementwise output

Here are the three kernels for a (4, 8) softmax, OPT=0:

void r_4(float* restrict data0, float* restrict data1) {
  for (int ridx0 = 0; ridx0 < 4; ridx0++) {
    float acc0;
    acc0 = -INFINITY;
    for (int ridx1 = 0; ridx1 < 8; ridx1++) {
      float val1 = acc0;
      float val2 = data0[((ridx0*8)+ridx1)];
      acc0 = ((val1>val2)?val1:val2);
    }
    float val3 = acc0;
    data1[ridx0] = val3;
  }
}

void r_4(float* restrict data0, float* restrict data1, float* restrict data2) {
  for (int ridx0 = 0; ridx0 < 4; ridx0++) {
    float acc0;
    acc0 = 0.0f;
    float val1 = data1[ridx0];
    for (int ridx1 = 0; ridx1 < 8; ridx1++) {
      float val2 = acc0;
      float val3 = data0[((ridx0*8)+ridx1)];
      acc0 = (val2+exp2f(((val3+(val1*-1.0f))*1.4426950408889634f)));
    }
    float val4 = acc0;
    data2[ridx0] = val4;
  }
}

void e_4_8(float* restrict data0, float* restrict data1, float* restrict data2, float* restrict data3) {
  for (int ridx0 = 0; ridx0 < 4; ridx0++) {
    float val1 = data1[ridx0];
    float val2 = data2[ridx0];
    for (int ridx1 = 0; ridx1 < 8; ridx1++) {
      int alu3 = ((ridx0*8)+ridx1);
      float val4 = data0[alu3];
      data3[alu3] = (exp2f(((val4+(val1*-1.0f))*1.4426950408889634f))*(1.0f/val2));
    }
  }
}

Notice that exp2f(...) appears in both the second and the third kernel. The exponentials are not materialized; they are recomputed in the consumer, because the kernelizer only cuts at reductions and roots. Recomputing one exp2f per element is cheaper than writing and re-reading a temporary, so here fusion wins. When the recomputation is a whole matmul instead of one exp2f, the answer flips, and that trade-off is the subject of Chapter 11.

9.9 argmax without an argmax op

To measure accuracy we need the index of the largest logit. The IR has MAX-reduce but no ARGMAX, and adding one would mean a reduction that carries two values (the best value and its index), touching rangeify and the renderer. Instead, a trick with ops we already have:

uopc/tensor.py · Tensor.argmax
  def argmax(self, axis=-1):
    """Index of the first maximum, as i32."""
    axis = axis % self.ndim
    n = self.shape[axis]
    idx = Tensor.arange(n).reshape((n,) + (1,) * (self.ndim - axis - 1))
    is_max = self.eq(self.max(axis, keepdim=True))
    return is_max.where(n - idx, 0).max(axis) * -1 + n

Step by step, for a row of length n:

  1. is_max is a 0/1 mask: where does the row equal its maximum?
  2. is_max.where(n - idx, 0) puts n - idx at every maximal position and 0 elsewhere. Since idx < n, the maximal positions all get positive values, and the smallest index gets the largest value.
  3. MAX-reducing picks n - first_idx, and * -1 + n turns it back into first_idx.

So ties go to the first occurrence, matching NumPy:

>>> x = Tensor(np.array([[0.1, 3.0, 3.0, -1.0], [2.0, 0.0, 5.0, 5.5]], np.float32))
>>> x.argmax(-1).numpy(), x.argmax(-1).dtype
[1 3] dtypes.i32

and it compiles to two reductions and a tiny epilogue (OPT=0):

void r_2(float* restrict data0, float* restrict data1, int* restrict data2, int* restrict data3) {
  for (int ridx0 = 0; ridx0 < 2; ridx0++) {
    int acc0;
    acc0 = (-2147483647-1);
    float val1 = data1[ridx0];
    for (int ridx1 = 0; ridx1 < 4; ridx1++) {
      int val2 = acc0;
      float val3 = data0[((ridx0*4)+ridx1)];
      int val4 = data2[ridx1];
      acc0 = ((val2>(((val3!=val1)!=1)?((val4*-1)+4):0))?val2:(((val3!=val1)!=1)?((val4*-1)+4):0));
    }
    int val5 = acc0;
    data3[ridx0] = val5;
  }
}

The int accumulator starting at INT_MIN is the identity() of an integer MAX-reduce from rangeify, and (val3!=val1)!=1 is how eq renders: "not not-equal". arange(n) is a real 4-element BUFFER (data2) here; a slightly smarter frontend would make it a RANGE-derived value and save the load.

9.10 Gradients

The Tensor-level autodiff entry point is short because Chapter 8 did the work:

uopc/tensor.py · Tensor.backward
  def backward(self, params: list[Tensor]) -> Tensor:
    """Set p.grad = d(self)/d(p) for each p. Gradients are lazy Tensors like everything else."""
    assert self.shape == (), "backward needs a scalar"
    root = getattr(self, "_graph", self.uop)
    for p, g in zip(params, gradient(root, [p.uop for p in params])):
      p.grad = Tensor(g) if p.grad is None else p.grad + Tensor(g)
    return self

backward takes the parameter list explicitly, calls gradient(loss, [p.uop ...]), and stores each result as a lazy Tensor in p.grad. Nothing is computed. Gradients are accumulated (p.grad + g) if backward is called twice, as in PyTorch. Continuing the session from 9.1:

>>> l = Linear(3, 4)
realize: 1 kernels, schedule 1.0 ms, run 0.1 ms
>>> [p.shape for p in parameters(l)]
[(3, 4), (4,)]
>>> loss = l(x).square().mean(); loss.backward(parameters(l))
>>> l.weight.grad, l.bias.grad
<Tensor (3, 4) f32> <Tensor (4,) f32>
>>> l.bias.grad.numpy()
realize: 4 kernels, schedule 24.1 ms, run 0.2 ms
[ 0.65893304  1.2755529   0.39464036 -1.0922788 ]

(The one kernel during Linear(3, 4) is the zero-initialized bias being realized into a buffer.) The gradient tests compare every interesting op against central finite differences in float64, with the same numgrad helper you met in Chapter 8:

tests/test_w05_tensor.py · test_gradients
@pytest.mark.parametrize("name,fn,npfn", [
  ("tanh_sum", lambda t: t.tanh().sum(), lambda x: np.tanh(x).sum()),
  ("softmax_xent", lambda t: t.log_softmax(-1)[:, 1].sum(), lambda x: (x - np.log(np.exp(x).sum(-1, keepdims=True)))[:, 1].sum()),
  ("max", lambda t: (t.max(-1) * t.sum(0).max()).sum(), lambda x: (x.max(-1) * x.sum(0).max()).sum()),
  ("pad_shrink", lambda t: t.pad(((1, 1), (0, 2)))[1:3, 1:].square().sum(),
   lambda x: (np.pad(x, ((1, 1), (0, 2)))[1:3, 1:] ** 2).sum()),
  ("sqrt_recip", lambda t: (t.square() + 1).sqrt().reciprocal().sum(), lambda x: (1 / np.sqrt(x ** 2 + 1)).sum()),
  ("matmul", lambda t: (t @ t.T).sin().sum(), lambda x: np.sin(x @ x.T).sum()),
])
def test_gradients(name, fn, npfn):
  x = r(3, 4)
  t = Tensor(x, requires_grad=True)
  fn(t).backward([t])
  np.testing.assert_allclose(t.grad.numpy(), numgrad(npfn, x), rtol=2e-2, atol=2e-3)

9.11 Layers

nn.py is written against the Tensor API only; it never touches a UOp. A layer is a Python object that holds parameter Tensors (created with requires_grad=True) and has a __call__:

uopc/nn.py · Linear
class Linear:
  def __init__(self, n_in: int, n_out: int, bias: bool = True):
    bound = 1 / math.sqrt(n_in)
    self.weight = Tensor.uniform(n_in, n_out, low=-bound, high=bound, requires_grad=True)
    self.bias = Tensor.zeros(n_out).realize() if bias else None
    if self.bias is not None: self.bias.requires_grad = True
  def __call__(self, x: Tensor) -> Tensor: return x.linear(self.weight, self.bias)

The weight is stored as (n_in, n_out) so that x @ w needs no transpose in the user's head (the compiler will decide where a transpose is worth making; see Chapter 10). The initialization is PyTorch's default for linear layers: uniform in +-1/sqrt(n_in), which keeps the variance of activations roughly constant from layer to layer. The bias starts at zero and is realized immediately, so it is a BUFFER rather than an EXPAND-of-a-CONST; otherwise "the bias" would be a constant baked into every kernel and its gradient would have nowhere to go.

Conv2d is the same pattern around conv2d, LayerNorm normalizes the last axis ((x - mean) / sqrt(var + eps) * weight + bias), and Embedding is a one-hot matmul:

uopc/nn.py · Embedding
class Embedding:
  """A one-hot matmul: differentiable without a scatter op."""
  def __init__(self, vocab: int, dim: int):
    self.vocab, self.weight = vocab, Tensor.randn(vocab, dim) * 0.02
    self.weight = self.weight.realize(); self.weight.requires_grad = True
  def __call__(self, idx: Tensor) -> Tensor: return idx.one_hot(self.vocab) @ self.weight

A real embedding layer is a gather (weight[idx]), and its gradient is a scatter-add. Our IR has neither, and autodiff refuses data-dependent INDEX. The one-hot product computes the same thing with ops autodiff understands: one_hot(idx) @ W selects rows, and its gradient one_hot^T @ dY scatters into them. It costs vocab times more FLOPs than a gather, which is fine for a character-level GPT (Chapter 14) and terrible for a 50,000-token vocabulary.

Finally, the optimizer needs a list of parameters. Rather than make every layer register its parameters, parameters() walks Python attributes, lists and dicts and collects every Tensor with requires_grad:

uopc/nn.py · parameters
def parameters(obj) -> list[Tensor]:
  """Every Tensor with requires_grad reachable through attributes, lists and dicts."""
  out, seen = [], set()
  def walk(o):
    if id(o) in seen: return
    seen.add(id(o))
    if isinstance(o, Tensor):
      if o.requires_grad: out.append(o)
    elif isinstance(o, (list, tuple)): [walk(x) for x in o]
    elif isinstance(o, dict): [walk(x) for x in o.values()]
    elif hasattr(o, "__dict__"): [walk(x) for x in vars(o).values()]
  walk(obj)
  return out

The seen set handles shared sub-objects (and cycles). The order is the attribute order of each object, which is deterministic, so the optimizer's state lines up with the same parameters every run.

9.12 Optimizers

An optimizer turns gradients into parameter updates. Every update in this library is built as a graph and realized in one batch, together with any extra tensors the caller wants (the loss).

SGD with momentum

Plain stochastic gradient descent is p <- p - lr * g. Momentum keeps a velocity v <- mu * v + g and steps along it, p <- p - lr * v, which smooths noisy gradients and speeds progress along consistent directions:

uopc/nn.py · SGD
class SGD(Optimizer):
  def __init__(self, params, lr=0.01, momentum=0.0):
    super().__init__(params)
    self.lr, self.momentum = lr, momentum
    self.v = [Tensor.zeros(*p.shape).realize() for p in params] if momentum else []
  def step(self, *extra: Tensor):
    """Update every parameter. Extra tensors (e.g. the loss) are realized in the same batch."""
    news = []
    for i, p in enumerate(self.params):
      g = p.grad
      if self.momentum:
        self.v[i].assign(self.v[i] * self.momentum + g); g = self.v[i]
      news.append(p.assign(p.detach() - g * self.lr))
    Tensor.realize_all(*news, *self.v, *extra)

The test checks the arithmetic by hand for the loss sum(w^2) (gradient 2w):

tests/test_w05_tensor.py · test_sgd_momentum_step
def test_sgd_momentum_step():
  w = Tensor(np.array([1.0, -2.0], np.float32), requires_grad=True)
  opt = SGD([w], lr=0.1, momentum=0.9)
  for _ in range(2):
    loss = w.square().sum()
    opt.zero_grad(); loss.backward(opt.params); opt.step()
  # v1 = 2w0, w1 = w0 - .1*v1 = .8 w0;  v2 = .9 v1 + 2 w1 = 3.4 w0, w2 = .8 w0 - .34 w0 = .46 w0
  np.testing.assert_allclose(w.numpy(), np.array([0.46, -0.92]), rtol=1e-5)

AdamW

Adam keeps two running averages per parameter: the first moment m (mean of the gradient) and the second moment v (mean of its square). With t counting steps from 1:

  • m <- b1 * m + (1 - b1) * g
  • v <- b2 * v + (1 - b2) * g^2
  • m_hat = m / (1 - b1^t), v_hat = v / (1 - b2^t) (bias correction: both start at 0, so early averages are too small; dividing by 1 - b^t undoes that exactly for a constant gradient)
  • p <- p * (1 - lr * wd) - lr * m_hat / (sqrt(v_hat) + eps)

Dividing by sqrt(v_hat) gives every parameter its own step size: parameters with large, noisy gradients take smaller steps. The (1 - lr * wd) factor is decoupled weight decay, the "W" in AdamW: it shrinks weights directly instead of adding wd * p to the gradient, which in Adam would get divided by sqrt(v_hat) and lose its meaning.

uopc/nn.py · AdamW
class AdamW(Optimizer):
  def __init__(self, params, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8, wd=0.0):
    super().__init__(params)
    self.lr, self.b1, self.b2, self.eps, self.wd, self.t = lr, b1, b2, eps, wd, 0
    self.m = [Tensor.zeros(*p.shape).realize() for p in params]
    self.v = [Tensor.zeros(*p.shape).realize() for p in params]
  def step(self, *extra: Tensor):
    self.t += 1
    # bias corrections are python floats baked into the graph as constants: each step's
    # graph differs only in two CONSTs, so we pass them as tensors to keep the JIT cache warm
    c1 = Tensor([self.lr / (1 - self.b1 ** self.t)])
    c2 = Tensor([1 / (1 - self.b2 ** self.t)])
    news = []
    for p, m, v in zip(self.params, self.m, self.v):
      g = p.grad
      m.assign(m * self.b1 + g * (1 - self.b1))
      v.assign(v * self.b2 + g.square() * (1 - self.b2))
      upd = m * c1.reshape(*([1] * p.ndim)) / ((v * c2.reshape(*([1] * p.ndim))).sqrt() + self.eps)
      news.append(p.assign(p.detach() * (1 - self.lr * self.wd) - upd))
    Tensor.realize_all(*news, *self.m, *self.v, *extra)

The only surprising lines are c1 and c2. The bias corrections depend on t, which changes every step. Written as Python floats they would become CONST nodes in the graph, the graph would be structurally different every step, and the structural JIT cache would miss every time. Passed as one-element tensors they are BUFFERs, whose contents do not participate in the cache key (see section 7.8 in Chapter 7). Here is the difference, measured by a scratch subclass AdamWFloats whose step is identical except that c1, c2 are floats (time per training step of the MNIST MLP, in ms, first eight steps):

AdamW        plans cached  5   step ms: 409 14 16 16 12 11 11 13
AdamWFloats  plans cached 12   step ms: 228 169 187 190 212 179 164 188

Same math, same results, more than ten times slower: every step misses the plan cache, so it is re-kernelized, re-rangeified and re-rendered in Python. (In this run gcc did not have to recompile: the .so files for these exact constants were already on disk from an earlier run of the same script. On a truly cold cache each step also pays gcc for the update kernels, and an earlier run measured 600-800 ms per step.) This is pitfall 17 in book/NOTES.md, and the general lesson is: anything that changes between iterations must be data, not structure.

9.13 Training MNIST

MNIST is 70,000 grayscale 28x28 images of handwritten digits (60,000 train, 10,000 test). The loader downloads the four gzip files once into ~/.cache/uopc/data and returns flat float32 images scaled to [0, 1] and int32 labels:

examples/datasets.py · mnist
def mnist():
  """(X_train (60000, 784) f32 in [0,1], Y_train (60000,) i32, X_test, Y_test)"""
  def images(n): return np.frombuffer(gzip.open(fetch(MNIST_URL + n, n)).read(), np.uint8, offset=16).reshape(-1, 784).astype(np.float32) / 255
  def labels(n): return np.frombuffer(gzip.open(fetch(MNIST_URL + n, n)).read(), np.uint8, offset=8).astype(np.int32)
  return (images("train-images-idx3-ubyte.gz"), labels("train-labels-idx1-ubyte.gz"),
          images("t10k-images-idx3-ubyte.gz"), labels("t10k-labels-idx1-ubyte.gz"))

The model is a two-layer perceptron, 784 inputs, 128 hidden units with ReLU, 10 outputs, trained with AdamW on random minibatches of 128:

examples/mnist_mlp.py
"""Week 5 milestone: train an MLP on MNIST with our own compiler.

  python examples/mnist_mlp.py            # ~94% test accuracy in 600 steps (~15 ms/step)
  DEBUG=2 python examples/mnist_mlp.py    # print every kernel with its time
"""
import sys, time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent))
import numpy as np
from datasets import mnist
from uopc.tensor import Tensor
from uopc.nn import Linear, AdamW, parameters


class MLP:
  def __init__(self):
    self.l1, self.l2 = Linear(784, 128), Linear(128, 10)
  def __call__(self, x: Tensor) -> Tensor: return self.l2(self.l1(x).relu())


def main(steps=600, bs=128):
  Xtr, Ytr, Xte, Yte = mnist()
  Tensor.manual_seed(0)
  model = MLP()
  opt = AdamW(parameters(model), lr=1e-3)
  rng = np.random.default_rng(0)
  t0 = time.perf_counter()
  for step in range(steps):
    idx = rng.integers(0, len(Xtr), bs)
    x, y = Tensor(Xtr[idx]), Tensor(Ytr[idx])
    loss = model(x).cross_entropy(y)
    opt.zero_grad()
    loss.backward(opt.params)
    opt.step(loss)
    if step % 100 == 0 or step == steps - 1:
      print(f"step {step:4d}  loss {loss.item():.4f}  {(time.perf_counter() - t0) / (step + 1) * 1e3:.1f} ms/step")
  pred = np.concatenate([model(Tensor(Xte[i:i + 1000])).argmax(-1).numpy() for i in range(0, len(Xte), 1000)])
  acc = (pred == Yte).mean()
  print(f"test accuracy {acc * 100:.2f}%")
  return acc


if __name__ == "__main__": main()

A training step is five lines: build the forward graph (model(x).cross_entropy(y)), clear the old gradients, build the backward graph, and let the optimizer realize the new parameters, moments and loss in one batch. Evaluation runs the model on the test set in chunks of 1000 and compares argmax with the labels. The real output, from a run with the current code (book/outputs/mnist.txt holds an earlier capture with identical losses and slightly slower steps: 460.7 ms for step 0 and 14.4 ms/step at the end):

step    0  loss 2.3134  284.8 ms/step
step  100  loss 0.3881  14.2 ms/step
step  200  loss 0.3272  13.1 ms/step
step  300  loss 0.2496  12.4 ms/step
step  400  loss 0.2643  11.9 ms/step
step  500  loss 0.2517  11.6 ms/step
step  599  loss 0.2162  11.6 ms/step
test accuracy 94.06%

(The ms/step column is a running average since step 0, so the expensive first step inflates the early rows.) The initial loss 2.31 is about ln 10 = 2.30, the loss of a uniform guess over ten classes, which is a good sanity check that initialization is sane. After 600 steps (about 1.3 epochs) the test accuracy is 94.06%. A longer run, a learning-rate schedule or a CNN would do better; the milestone is that a model written in your compiler learns.

Step 0 versus the rest. The first step takes about 300 ms because it plans and renders 32 kernels (and, on a cold cache, runs gcc 32 times: with an empty UOPC_CACHE the same step took 3.7 s). Every later step has the same graph structure: same ops, same shapes, different buffers. _structure_key (Chapter 7) maps it to the cached plan, so the steady-state step only builds graphs, hashes them, and calls 32 already-loaded C functions.

9.14 Anatomy of one training step

DEBUG=2 prints every kernel with its wall time and a rough GFLOP/s. Here is the second training step (the first warm one) of the MLP above, from a run made while this chapter was being written (a rerun with the current code produces the same 32 kernels in the same order; individual times move by tens of percent from run to run):

  e_1_128_784                       0.279 ms     0.00 GFLOP/s
  r_128_128                         0.431 ms    59.65 GFLOP/s
  r_128                             0.047 ms     0.19 GFLOP/s
  e_128_128                         0.027 ms     1.84 GFLOP/s
  e_1_10_128                        0.013 ms     0.00 GFLOP/s
  r_128_10                          0.180 ms     1.82 GFLOP/s
  r_128                             0.051 ms     0.05 GFLOP/s
  r_128                             0.030 ms     0.26 GFLOP/s
  r_128                             0.041 ms     0.65 GFLOP/s
  r_128                             0.015 ms     0.43 GFLOP/s
  e_128_10                          0.048 ms     0.75 GFLOP/s
  r_128_128                         0.132 ms     2.48 GFLOP/s
  e_128_128                         0.052 ms     0.95 GFLOP/s
  r_128_784                         0.626 ms    41.01 GFLOP/s
  e_784_128                         0.375 ms     0.80 GFLOP/s
  e_784_128                         0.484 ms     0.83 GFLOP/s
  e_784_128                         0.597 ms     1.51 GFLOP/s
  r_128                             0.052 ms     0.32 GFLOP/s
  e_128                             0.035 ms     0.01 GFLOP/s
  e_128                             0.015 ms     0.03 GFLOP/s
  e_128                             0.025 ms     0.05 GFLOP/s
  r_10_128                          0.187 ms     1.75 GFLOP/s
  e_128_10                          0.038 ms     0.10 GFLOP/s
  e_128_10                          0.017 ms     0.30 GFLOP/s
  e_128_10                          0.029 ms     0.40 GFLOP/s
  r_10                              0.010 ms     0.13 GFLOP/s
  e_10                              0.014 ms     0.00 GFLOP/s
  e_10                              0.013 ms     0.00 GFLOP/s
  e_10                              0.023 ms     0.00 GFLOP/s
  r_128                             0.038 ms     0.41 GFLOP/s
  r_1                               0.010 ms     0.01 GFLOP/s
  e_1                               0.015 ms     0.00 GFLOP/s
realize: 32 kernels, schedule 1.6 ms, run 4.4 ms

Kernel names encode the output shape: r_ kernels contain a reduction, e_ kernels are purely elementwise. By reading the C (DEBUG=4) you can assign every kernel to a role:

The 32 kernels of one MNIST training step A strip of the 32 kernels of one warm training step of the MNIST MLP, in execution order, each as wide as its measured run time (3.95 ms in total). Forward kernels take 1.01 ms, backward 1.21 ms, the twelve AdamW update kernels 1.67 ms and the loss 0.06 ms. The six kernels that touch the 784 by 128 weight W1 take 2.79 ms. ONE WARM STEP · 32 KERNELS IN EXECUTION ORDER · WIDTH = RUN TIME W1ᵀ copy x @ W1 h@W2 dh dW1 new m new v new W1 dW2 W1 (784x128): 2.79 of 3.95 ms 0 ms 1 ms 2 ms 3 ms forward 1.01 ms backward 1.21 ms AdamW update 1.67 ms loss 0.06 ms the optimizer costs more than the forward pass

Summed, the kernels take 3.95 ms; the run 4.4 ms in the realize line also includes the Python and ctypes cost of 32 calls. The six kernels that touch W1, the 784x128 matrix that holds 99% of the parameters, account for 2.79 ms of it, and the twelve optimizer kernels together cost more than the whole forward pass.

Forward (7 kernels). e_1_128_784 copies W1 into a transposed temporary and r_128_128 is x @ W1; the copy is the layout pass from Chapter 10, which makes the reduce axis contiguous for both operands. e_128_128 is relu(h + b1), materialized by stage_broadcast_compute (Chapter 7) because it is broadcast into the next matmul and would otherwise be recomputed for every output column. e_1_10_128 and r_128_10 are the copy-and-matmul pair for W2. Two r_128 kernels compute the row max and the sum of exponentials of log_softmax.

Backward (10 kernels). Four small kernels for the gradient of cross-entropy with respect to the logits (the last of them, e_128_10, writes dL/dlogits), r_128_128 and e_128_128 for dL/dh through W2 and the ReLU, r_128_784 for dW1, r_10_128 for dW2, and two small reductions over the batch for the bias gradients. The scheduler is free to interleave phases: the third kernel of the step (r_128 reading the labels) is already a backward kernel, because it depends only on the labels. The ReLU backward kernel is a one-liner:

void e_128_128(float* restrict data0, float* restrict data1, float* restrict data2, float* restrict data3) {
  for (int ridx0 = 0; ridx0 < 128; ridx0++) {
    for (int ridx1 = 0; ridx1 < 128; ridx1++) {
      int alu1 = ((ridx0*128)+ridx1);
      float val2 = data0[alu1];
      float val3 = data1[ridx1];
      float val4 = data2[alu1];
      data3[alu1] = ((0.0f<(val2+val3))?val4:0.0f);
    }
  }
}

It recomputes h + b1 > 0 from the pre-activation instead of storing a mask: that is the WHERE rule from autodiff.py (s[0].where(g, 0)) fused with the forward expression that produced the condition.

Update (12 kernels). Three per parameter: new m, new v, new p. They are separate kernels because m, v and p are all roots of the realize_all call, and roots are materialization points. The W1 update reads the freshly written m and v:

void e_784_128(float* restrict data0, float* restrict data1, float* restrict data2, float* restrict data3, float* restrict data4, float* restrict data5) {
  float val1 = data2[0];
  float val2 = data4[0];
  #pragma omp parallel for
  for (int ridx0 = 0; ridx0 < 784; ridx0++) {
    for (int ridx1 = 0; ridx1 < 128; ridx1++) {
      int alu3 = ((ridx0*128)+ridx1);
      float val4 = data0[alu3];
      float val5 = data1[alu3];
      float val6 = data3[alu3];
      data5[alu3] = (val4+(((val5*val1)*(1.0f/(sqrtf((val6*val2))+1e-08f)))*-1.0f));
    }
  }
}

data2[0] and data4[0] are c1 and c2: one-element BUFFERs, loaded once, outside the loop. With wd = 0 the decay factor (1 - lr * wd) folded to 1 and vanished.

Loss (3 kernels). Per-example cross-entropy, the sum over the batch, and the final scale by -1/128.

Two observations are worth making before moving on. First, two of the backward kernels (a sum over classes and a tie count) and part of a third compute the gradient that flows through the max in log_softmax. Mathematically that gradient is exactly zero: shifting all logits by a constant does not change log_softmax, so the derivative with respect to the shift vanishes. Autodiff does not know that, and faithfully computes a zero the expensive way. Second, the m and v updates for W1 read the gradient at data1[((ridx1*784)+ridx0)]: the gradient is produced as (128, 784) (that is dW1 transposed) and read with a stride of 784. Both are opportunities, not bugs.

9.15 What is slow, and why

The steady-state step costs about 12 ms (the last line of the run in 9.13), and the 32 kernels account for roughly 4 to 6 ms of it. Where does the rest go? A scratch script that times the three phases of 50 warm steps (graph building for the forward pass, backward, and opt.step) printed:

forward graph                  1.87 ms/step
backward graph                 1.71 ms/step
step (structure key+run)       9.81 ms/step

and cProfile over 50 warm steps (the profiler inflates everything, but the ranking is what matters):

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     1600    0.168    0.000    0.198    0.000 uopc/runtime.py:53(__call__)
      150    0.080    0.001    0.157    0.001 uopc/uop.py:131(toposort)
62950/32100    0.076    0.000    0.126    0.000 uopc/uop.py:100(_argkey)
    11750    0.075    0.000    0.125    0.000 uopc/uop.py:230(_infer)
    18650    0.050    0.000    0.326    0.000 uopc/uop.py:111(__new__)
       50    0.037    0.001    0.190    0.004 uopc/kernelize.py:180(_structure_key)

Program.__call__ is the 32 C calls per step (the kernels themselves plus ctypes argument marshalling). Everything else is Python: about 370 UOp constructions per step (18,650 over 50 steps) (each one hash-consed and type-checked by _infer), topologically sorting the step's graph, and hashing its structure for the plan cache. This is the tax of a Python compiler that rebuilds the graph every step. NOTES 22 tells how it was once much worse: UOp.__eq__/__hash__ written in Python, and Enum.__hash__ being a Python function, made the GPT compile take 163 s for 12 steps until identity hashing (in C) cut it to 47 s. Production systems avoid the rebuild entirely by capturing the step once and replaying it (tinygrad's TinyJit, CUDA graphs).

What about the kernels themselves? The two big matmuls run at 41-60 GFLOP/s, the elementwise kernels at around 1 GFLOP/s or less. For elementwise work that is expected: they are bound by memory and by fixed per-call overhead, as Chapter 10 will quantify with the roofline model. Running the whole training with the optimizer switched off (OPT=0, book/outputs/mnist_opt0.txt) gives the same losses digit for digit and 17.2 ms/step instead of 14.4 (both from the earlier captures): on a model this small the Python overhead hides most of the kernel speedup. On the 1024x1024 GEMM of the next chapter it is the other way around, and the optimizer is worth a factor of more than 50.

9.16 Building it yourself: the week-5 milestone

If you are implementing this for the course, here is an order that keeps every step testable:

  1. Tensor.__init__, numpy(), realize() on top of your realize from Chapter 7. Test: Tensor(a).numpy() == a.
  2. _broadcast_shape, broadcast_to, _binop, and the arithmetic dunders. Test against NumPy on random shapes, including (2, 3, 4) + (4,), (2, 1, 4) * (3, 1) and scalar operands on both sides.
  3. Movement methods and __getitem__. Test chains of them against NumPy (test_ops[getitem_cat]).
  4. _reduce with arbitrary axes and keepdim; mean, var, max, min. Test sum((0, 2)).
  5. matmul, then softmax/log_softmax (check softmax(x + 1000) does not produce NaN), one_hot, cross_entropy, argmax (check ties).
  6. backward, then gradient-check every op you wrote with finite differences.
  7. Linear, parameters, SGD; overfit 64 random points (test_training_reduces_loss). Only then AdamW and the real dataset.
  8. MNIST. Watch the first loss (should be near 2.30) and the ms/step after step 0 (should drop by more than an order of magnitude; if it does not, your JIT cache is missing).

Pitfalls

A changing Python constant defeats the JIT cache (NOTES 17). Section 9.12 measured it: AdamW's bias corrections as floats made every step more than ten times slower. Learning-rate schedules, dropout masks and step counters must all be tensors (see Chapter 7, Pitfalls).

Forgetting to realize optimizer state. realize_all(*news, *self.m, *self.v, *extra) is not decoration. A scratch LeakyAdamW that realizes only the parameters still trains correctly, because the new m and v get computed inside the parameter kernels. But m and v themselves stay lazy, so each step's graph contains the previous step's graph:

step 0:     385 ms   graph behind m[0]:  134 nodes
step 1:     263 ms   graph behind m[0]:  235 nodes
step 2:     517 ms   graph behind m[0]:  336 nodes
step 3:     616 ms   graph behind m[0]:  437 nodes
step 4:     771 ms   graph behind m[0]:  538 nodes
step 5:    1102 ms   graph behind m[0]:  639 nodes

The graph grows by 101 nodes per step, the structure never repeats, and every step recompiles a bigger program. Any state that outlives a step must be realized every step.

Python scalars took the tensor's dtype. The first version of _binop built every Python scalar with self.dtype, so Tensor([1, 2, 3]) * 0.5 multiplied by int(0.5) == 0, and Tensor([1, 2, 3]).mean() was 0: NumPy promotes to float, we silently truncated. The fix is the promotion line at the top of _binop (an int tensor meeting a Python float is cast to f32) and a mean that always sums in f32. Now:

>>> t = Tensor([1, 2, 3])
>>> (t * 0.5).numpy(), t.mean().item()
[0.5 1.  1.5] 2.0
tests/test_w05_tensor.py · test_python_float_promotes_int_tensors
def test_python_float_promotes_int_tensors():
  t = Tensor(np.array([1, 2, 3], np.int32))
  np.testing.assert_allclose((t * 0.5).numpy(), [0.5, 1.0, 1.5])
  assert t.mean().item() == 2.0

A kernel whose input simplified away overwrote that input. Found while writing this chapter. In make_plan, a kernel's argument list (refs) was every BUFFER of its shaped body, collected before lowering. The renderer, however, emits a C parameter only for buffers the simplified program still reads. When simplification deletes a read entirely, the two lists disagree and every pointer shifts by one. t * 0 on an integer tensor folds to the constant 0, the kernel's signature becomes void e_3(int* restrict data1), and ctypes passed t's own buffer as data1: the first version printed [0 0 0] [0 0 0] for (t * 0).numpy(), t.numpy(), writing the "output" into the input. (Floats escape, because x * 0.0 is not folded: it is not 0.0 for NaN or infinity.) The fix keeps only the slots that survive in the lowered program:

uopc/kernelize.py:159-163
    refs = [ph_ref[b] if b.arg.slot < 0 else ("in", in_idx[b]) for b in bufs] + [ref[u]]
    # simplification can delete a read (x*0 -> 0): the C signature then has fewer parameters,
    # so bind exactly the slots that survived, or every later argument shifts by one
    alive = sorted({b.arg.slot for b in ast.toposort() if b.op is Ops.BUFFER})
    refs = [refs[s] for s in alive]

The generated kernel still has one parameter, but it is now bound to the right buffer:

void e_3(int* restrict data1) {
  for (int ridx0 = 0; ridx0 < 3; ridx0++) {
    data1[ridx0] = 0;
  }
}

>>> (t * 0).numpy(), t.numpy()
[0 0 0] [1 2 3]
tests/test_w05_tensor.py · test_deleted_reads_do_not_shift_kernel_arguments
def test_deleted_reads_do_not_shift_kernel_arguments():
  t = Tensor(np.array([1, 2, 3], np.int32))
  assert (t * 0).numpy().tolist() == [0, 0, 0]          # simplifies to a constant: t is no longer read...
  assert t.numpy().tolist() == [1, 2, 3]                # ...and must not be overwritten

The general lesson: whenever two passes each compute "the list of parameters" independently, they will eventually disagree. Derive one from the other.

Unstable softmax. If you write x.exp() / x.exp().sum(-1, keepdim=True), it works on small random inputs and returns NaN on real logits above about 88. Always subtract the max; test with x + 1000.

Fusion can recompute a lot (NOTES 15). In softmax the recomputed value is one cheap exp2f, but in a transformer a LayerNorm output feeding a matmul is recomputed for every output column (see Chapter 7, Pitfalls). If one kernel in DEBUG=2 output is mysteriously slow, read its C and look for an expensive expression inside the reduce loop.

Reading the loss before backward used to zero the gradients. .item() realizes the loss, and realizing rebinds loss.uop to the BUFFER holding the result (9.1). In the first version that was the end of the story: the graph that connected the loss to the parameters was gone, gradient() found no path and returned zeros, with no error and no warning:

>>> w = Tensor(np.array([1.0, -2.0], np.float32), requires_grad=True)
>>> loss = w.square().sum()
>>> loss.item()
5.0
>>> loss.backward([w]); w.grad.numpy()
[0. 0.]

The fix is the t._graph = t.uop line in realize_all (9.2): a realized tensor remembers the graph it came from, and backward differentiates that graph if it exists:

uopc/tensor.py · Tensor.backward
  def backward(self, params: list[Tensor]) -> Tensor:
    """Set p.grad = d(self)/d(p) for each p. Gradients are lazy Tensors like everything else."""
    assert self.shape == (), "backward needs a scalar"
    root = getattr(self, "_graph", self.uop)
    for p, g in zip(params, gradient(root, [p.uop for p in params])):
      p.grad = Tensor(g) if p.grad is None else p.grad + Tensor(g)
    return self

The same session now prints 5.0 and then [ 2. -4.]. It is still good practice to read the loss after opt.step(loss), as mnist_mlp.py does, so the loss is computed in the same batch of kernels.

detach() used to be a no-op for gradients. The first detach returned a new Tensor around the same UOp, and gradient() follows UOps, not Tensor objects, so the gradient flowed straight through. Now detach wraps the value in a marker node:

uopc/tensor.py · Tensor.detach
  def detach(self) -> Tensor:
    """Same value, but gradients stop here. (Kernelize strips the marker, so it costs nothing to run.)"""
    return Tensor(UOp(Ops.STAGE, (self.uop,), "detach"))

Autodiff's rule for STAGE passes the gradient through unchanged, except when the STAGE's arg is "detach", where it returns no gradient. The kernelizer removes the marker before anything else, so it costs nothing at run time:

uopc/kernelize.py · strip_detach
strip_detach = PatternMatcher([(UPat(Ops.STAGE, arg="detach", src=(UPat.var("x"),)), lambda x: x)])

Both fixes are pinned by one test: realizing before backward keeps the gradient 2w, and (w.detach() * w).sum() has gradient w, not 2w:

tests/test_w05_tensor.py · test_item_before_backward_and_detach
def test_item_before_backward_and_detach():
  w = Tensor(np.array([1.0, 2.0], np.float32), requires_grad=True)
  loss = (w * w).sum()
  loss.item()                                          # realizing the loss must keep its graph
  loss.backward([w])
  np.testing.assert_allclose(w.grad.numpy(), [2.0, 4.0])
  w.grad = None
  (w.detach() * w).sum().backward([w])                 # d/dw (c * w) with c = w held constant
  np.testing.assert_allclose(w.grad.numpy(), [1.0, 2.0])

Exercises

  1. ★ Count kernels. Predict how many kernels (x - x.mean(-1, keepdim=True)).square().sum() produces for x of shape (64, 64), then check with make_plan. Hint: count REDUCEs.
  2. ★ Predict argmax. What does Tensor(np.array([[3., 1., 3.]], np.float32)).argmax(-1).numpy() print, and why? What would happen if the trick used idx instead of n - idx?
  3. ★ Std and cross-entropy. Implement Tensor.std(axis, keepdim) and Tensor.logsumexp-based cross_entropy (use logsumexp instead of log_softmax). Gradient-check both.
  4. ★★ A differentiable gather. Implement Tensor.gather(axis, idx) for integer index tensors using one_hot and a reduce, so that it is differentiable with respect to the data. What does it cost in FLOPs compared to a real gather?
  5. ★★ A learning-rate schedule. Add a learning-rate schedule (linear warmup, then cosine decay) to AdamW without causing JIT cache misses. Verify with DEBUG=1 that schedule stays around a millisecond after step 0.
  6. ★★ Unused buffers. Revert the alive filter in make_plan (in a scratch copy), watch test_deleted_reads_do_not_shift_kernel_arguments fail, then fix it the other way: render every buffer of the body as a C parameter even when it is unused. Which fix is better, and why? Hint: think about the structural cache and about identical kernels rendering identical C.
  7. ★★ Dropout. Implement dropout: x.dropout(p) in training mode multiplies by a random 0/(1/(1-p)) mask. Where must the randomness come from so that the graph structure is the same every step?
  8. ★★ A small CNN. Train a small CNN on MNIST (Conv2d(1, 8, 3), ReLU, max_pool2d, Linear). Report accuracy and ms/step, and identify the slowest kernel with DEBUG=2. Explain its cost from its C source.
  9. ★★ Detach the max. Use detach() to stop the gradient through the max in log_softmax (it is mathematically zero, see section 9.14). Count the kernels per training step before and after, and check that the losses are unchanged to the printed precision.
  10. ★★★ Fast convolution. Make conv2d fast: instead of a STACK of shrinks selected with a WHERE chain, express the kernel taps as extra reduce axes (shape (N, O, OH, OW, C, KH, KW) reduced over the last three). Compare the generated C and the speed against section 9.7's kernel on a (64, 16, 28, 28) input.
Solutions to ★ exercises

1. Two: make_plan reports ['r_64', 'r_1']. The mean REDUCE is one materialization point (r_64; its * (1/64) is elementwise and fuses into its consumer). The subtraction and the square are elementwise and fuse into the final sum(), which is both a REDUCE and the root, so it is a single point (r_1). Count materialization points, and remember that a root which is itself a REDUCE counts once.

2. [0]. The row max is 3, is_max = [1, 0, 1], where(is_max, 3 - idx, 0) = [3, 0, 1], max is 3, and 3 * -1 + 3 = 0. With idx instead of n - idx the max would pick the last maximal index (2), and the zero for non-maximal positions would collide with index 0.

3. std = var.sqrt(). For cross-entropy: (self.logsumexp(-1).reshape(-1) - (self * labels.one_hot(n)).sum(-1)).mean(). Both match finite differences to about 5e-8. (With the first version of realize_all, reading the loss with .item() before backward made this check fail with zero gradients; see the Pitfalls.)

Quiz

  1. Why is broadcasting free in this compiler, while in an eager library np.broadcast_to(b, shape).copy() costs memory traffic?
  2. REDUCE only reduces the last n axes. How does x.sum(0) on a (4, 5) tensor become a valid graph?
  3. Write matmul for (M, K) @ (K, N) as a sequence of movement ops, one MUL and one REDUCE, with the shape after each step.
  4. Why does log_softmax subtract the row max, and why does doing so not change the result?
  5. Why are AdamW's bias corrections passed as one-element tensors rather than Python floats?
  6. How many kernels does a softmax over the last axis produce, and which values are recomputed rather than stored?
  7. Why does Embedding use a one-hot matmul rather than indexing?
  8. After y.numpy(), what is y.uop, and why does that matter for graphs built on y later?
Answers
  1. EXPAND is lowered to an index expression that ignores the broadcast loop (index 0 on the stretched axis); no data is copied, the same element is simply re-read.
  2. _reduce permutes the reduced axis to the end (permute(1, 0), shape (5, 4)) and REDUCEs 1 axis.
  3. A.reshape(M, 1, K).expand(M, N, K), B.permute(1, 0).reshape(1, N, K).expand(M, N, K), MUL gives (M, N, K), REDUCE(ADD, 1) gives (M, N).
  4. To keep every exponent at most 0 (no overflow) with the largest term exactly 1 (no 0/0). The constant cancels between numerator and denominator (or, in log space, between the two terms).
  5. They change every step; as CONSTs they would change the graph structure and miss the JIT cache, so every step would recompile.
  6. Three: max, sum of exponentials, and the elementwise output. The exponentials are recomputed in the sum kernel and again in the output kernel.
  7. Autodiff cannot differentiate a data-dependent INDEX (its gradient is a scatter). One-hot times weights is an ordinary broadcast-multiply-reduce whose gradient is another matmul.
  8. A BUFFER (or a reshape of one) holding the result. Later graphs start from that data instead of re-including, and recomputing, the whole chain that produced it.

Further reading

  • NumPy documentation, "Broadcasting" (numpy.org/doc/stable/user/basics.broadcasting.html): the rules this chapter implements.
  • D. Kingma and J. Ba, "Adam: A Method for Stochastic Optimization", ICLR 2015; I. Loshchilov and F. Hutter, "Decoupled Weight Decay Regularization", ICLR 2019 (AdamW).
  • K. Chellapilla, S. Puri, P. Simard, "High Performance Convolutional Neural Networks for Document Processing", 2006: the im2col formulation of convolution.
  • Y. LeCun, C. Cortes, C. Burges, "The MNIST database of handwritten digits".
  • tinygrad, tinygrad/tensor.py: a production-sized version of this file, with the same lazy design.
  • A. Karpathy, micrograd: the smallest possible autodiff engine, a good contrast to graph-level autodiff.

Full code for this chapter

Every file this chapter builds, exactly as it is in the repository. Click a file to expand it.

uopc/tensor.py · 333 lines
uopc/tensor.py
"""Week 5: a lazy Tensor frontend.

A Tensor is a thin wrapper around a shaped UOp graph. Nothing runs until you call
.realize() / .numpy(): then the graph is kernelized, rangeified, rendered to C,
compiled, and executed. Every op below is a few lines of graph building on top of
the course's primitives -- there is no other math library underneath.
"""
from __future__ import annotations
import math, os
import numpy as np
from uopc.uop import UOp, Ops, prod
from uopc.dtype import DType, dtypes
from uopc.autodiff import gradient
from uopc.composite import pad as _pad, flip_axis
from uopc.kernelize import new_buffer, buffer_data, base_buffer, realize as _realize
from uopc.runtime import np_dtype

LOG2E, LN2 = 1 / math.log(2), math.log(2)


def _broadcast_shape(*shapes):
  n = max(len(s) for s in shapes)
  shapes = [(1,) * (n - len(s)) + tuple(s) for s in shapes]
  out = []
  for dims in zip(*shapes):
    d = max(dims)
    if any(x not in (1, d) for x in dims): raise ValueError(f"can't broadcast {shapes}")
    out.append(d)
  return tuple(out)


class Tensor:
  rng = np.random.default_rng(0)
  training = False
  fuse_reduce = False
  opt = None   # week 6: kernel optimizer (set below)

  def __init__(self, data, dtype: DType | None = None, requires_grad: bool = False):
    if isinstance(data, UOp): self.uop = data
    else:
      arr = np.asarray(data)
      if dtype is None: dtype = dtypes.f32 if arr.dtype.kind == "f" else dtypes.bool if arr.dtype.kind == "b" else dtypes.i32
      arr = arr.astype(np_dtype(dtype))
      buf = new_buffer(dtype, arr.size, arr)
      self.uop = buf.reshape(arr.shape) if arr.shape != buf.shape else buf
    self.requires_grad = requires_grad
    self.grad: Tensor | None = None

  # ---- basics -------------------------------------------------------------
  @property
  def shape(self) -> tuple[int, ...]: return self.uop.shape
  @property
  def dtype(self) -> DType: return self.uop.dtype
  @property
  def ndim(self) -> int: return len(self.shape)
  def numel(self) -> int: return prod(self.shape)
  def __repr__(self): return f"<Tensor {self.shape} {self.dtype.name}{' grad' if self.requires_grad else ''}>"

  def realize(self) -> Tensor:
    Tensor.realize_all(self)
    return self

  @staticmethod
  def realize_all(*ts: Tensor):
    todo = [t for t in ts if t.uop.op is not Ops.BUFFER and not (t.uop.op is Ops.RESHAPE and base_buffer(t.uop).op is Ops.BUFFER and prod(t.shape) == base_buffer(t.uop).arg.shape[0])]
    if not todo: return
    bufs = _realize([t.uop for t in todo], Tensor.fuse_reduce, Tensor.opt)
    for t, b in zip(todo, bufs):
      t._graph = t.uop   # keep the graph: loss.item() before loss.backward() must still have gradients
      t.uop = b.reshape(t.shape) if b.shape != t.shape else b

  def numpy(self) -> np.ndarray:
    self.realize()
    return buffer_data(base_buffer(self.uop)).reshape(self.shape).copy()

  def item(self): return self.numpy().item()
  def detach(self) -> Tensor:
    """Same value, but gradients stop here. (Kernelize strips the marker, so it costs nothing to run.)"""
    return Tensor(UOp(Ops.STAGE, (self.uop,), "detach"))

  def assign(self, x: Tensor) -> Tensor:
    """Make this tensor (e.g. a parameter) hold x's value. Functional: rebinds to a new buffer."""
    assert x.shape == self.shape and x.dtype == self.dtype, f"assign {x.shape} to {self.shape}"
    self.uop = x.uop
    return self

  # ---- creation -----------------------------------------------------------
  @staticmethod
  def full(shape, v, dtype: DType = dtypes.f32) -> Tensor:
    c = UOp.const(v, dtype)
    return Tensor(c.reshape((1,) * len(shape)).expand(tuple(shape)) if shape else c)
  @staticmethod
  def zeros(*shape, dtype=dtypes.f32) -> Tensor: return Tensor.full(shape, 0, dtype)
  @staticmethod
  def ones(*shape, dtype=dtypes.f32) -> Tensor: return Tensor.full(shape, 1, dtype)
  @staticmethod
  def randn(*shape, requires_grad=False) -> Tensor:
    return Tensor(Tensor.rng.standard_normal(shape).astype(np.float32), requires_grad=requires_grad)
  @staticmethod
  def uniform(*shape, low=-1.0, high=1.0, requires_grad=False) -> Tensor:
    return Tensor(Tensor.rng.uniform(low, high, shape).astype(np.float32), requires_grad=requires_grad)
  @staticmethod
  def kaiming(*shape, requires_grad=True) -> Tensor:
    bound = 1 / math.sqrt(prod(shape[1:]) if len(shape) > 1 else shape[0])
    return Tensor.uniform(*shape, low=-bound, high=bound, requires_grad=requires_grad)
  @staticmethod
  def arange(n: int) -> Tensor: return Tensor(np.arange(n, dtype=np.int32))
  @staticmethod
  def manual_seed(s: int): Tensor.rng = np.random.default_rng(s)

  # ---- elementwise ----------------------------------------------------------
  def _const(self, x) -> Tensor: return x if isinstance(x, Tensor) else Tensor.full((), x, self.dtype)

  def broadcast_to(self, shape) -> Tensor:
    shape = tuple(shape)
    if self.shape == shape: return self
    u = self.uop
    if len(u.shape) < len(shape): u = u.reshape((1,) * (len(shape) - len(u.shape)) + u.shape)
    return Tensor(u.expand(shape) if u.shape != shape else u)

  def _binop(self, op: Ops, x, reverse=False) -> Tensor:
    # numpy-style promotion for python scalars: int tensor (op) 0.5 is a float computation
    if isinstance(x, float) and not self.dtype.is_float: return self.float()._binop(op, x, reverse)
    x = self._const(x)
    a, b = (x, self) if reverse else (self, x)
    shape = _broadcast_shape(a.shape, b.shape)
    return Tensor(UOp(op, (a.broadcast_to(shape).uop, b.broadcast_to(shape).uop)))

  def __add__(self, x): return self._binop(Ops.ADD, x)
  def __radd__(self, x): return self._binop(Ops.ADD, x, True)
  def __mul__(self, x): return self._binop(Ops.MUL, x)
  def __rmul__(self, x): return self._binop(Ops.MUL, x, True)
  def __neg__(self): return self * -1
  def __sub__(self, x): return self + -self._const(x)
  def __rsub__(self, x): return self._const(x) + -self
  def __truediv__(self, x): return self * self._const(x).reciprocal()
  def __rtruediv__(self, x): return self._const(x) * self.reciprocal()
  def __lt__(self, x): return self._binop(Ops.CMPLT, x)
  def __gt__(self, x): return self._binop(Ops.CMPLT, x, True)
  def __ge__(self, x): return (self < x).logical_not()
  def __le__(self, x): return (self > x).logical_not()
  def eq(self, x): return self._binop(Ops.CMPNE, x).logical_not()
  def ne(self, x): return self._binop(Ops.CMPNE, x)
  def logical_not(self): return self._binop(Ops.CMPNE, True)
  def maximum(self, x): return self._binop(Ops.MAX, x)
  def minimum(self, x): return -((-self).maximum(-self._const(x)))
  def where(self, a, b) -> Tensor:
    a = a if isinstance(a, Tensor) else Tensor.full((), a)
    b = b if isinstance(b, Tensor) else Tensor.full((), b, a.dtype)
    shape = _broadcast_shape(self.shape, a.shape, b.shape)
    return Tensor(UOp(Ops.WHERE, (self.broadcast_to(shape).uop, a.broadcast_to(shape).uop, b.broadcast_to(shape).uop)))
  def cast(self, d: DType) -> Tensor: return Tensor(self.uop.cast(d))
  def float(self) -> Tensor: return self.cast(dtypes.f32)

  def reciprocal(self): return Tensor(self.uop.recip())
  def exp2(self): return Tensor(UOp(Ops.EXP2, (self.uop,)))
  def log2(self): return Tensor(UOp(Ops.LOG2, (self.uop,)))
  def exp(self): return (self * LOG2E).exp2()
  def log(self): return self.log2() * LN2
  def sqrt(self): return Tensor(UOp(Ops.SQRT, (self.uop,)))
  def rsqrt(self): return self.sqrt().reciprocal()
  def sin(self): return Tensor(UOp(Ops.SIN, (self.uop,)))
  def cos(self): return (self + math.pi / 2).sin()
  def square(self): return self * self
  def abs(self): return self.maximum(-self)
  def relu(self): return (self > 0).where(self, 0.0)
  def sigmoid(self): return ((-self).exp() + 1).reciprocal()
  def tanh(self): return (self * 2).sigmoid() * 2 - 1
  def gelu(self): return 0.5 * self * (1 + (math.sqrt(2 / math.pi) * (self + 0.044715 * self * self * self)).tanh())
  def silu(self): return self * self.sigmoid()

  # ---- movement -------------------------------------------------------------
  def reshape(self, *shape) -> Tensor:
    shape = tuple(shape[0]) if len(shape) == 1 and isinstance(shape[0], (tuple, list)) else shape
    if -1 in shape:
      known = prod(s for s in shape if s != -1)
      shape = tuple(self.numel() // known if s == -1 else s for s in shape)
    return self if shape == self.shape else Tensor(self.uop.reshape(shape))
  def permute(self, *order) -> Tensor:
    order = tuple(order[0]) if len(order) == 1 and isinstance(order[0], (tuple, list)) else order
    order = tuple(o % self.ndim for o in order)
    return self if order == tuple(range(self.ndim)) else Tensor(self.uop.permute(order))
  def transpose(self, a=-1, b=-2) -> Tensor:
    order = list(range(self.ndim))
    order[a], order[b] = order[b], order[a]
    return self.permute(order)
  @property
  def T(self) -> Tensor: return self.transpose()
  def expand(self, *shape) -> Tensor:
    shape = tuple(shape[0]) if len(shape) == 1 and isinstance(shape[0], (tuple, list)) else shape
    return self.broadcast_to(tuple(s if s != -1 else o for s, o in zip(shape, self.shape)))
  def flatten(self, start=0) -> Tensor: return self.reshape(self.shape[:start] + (-1,))
  def unsqueeze(self, d) -> Tensor:
    d = d % (self.ndim + 1)
    return self.reshape(self.shape[:d] + (1,) + self.shape[d:])
  def shrink(self, arg) -> Tensor:
    arg = tuple((0, s) if a is None else a for a, s in zip(arg, self.shape))
    return self if all(a == (0, s) for a, s in zip(arg, self.shape)) else Tensor(self.uop.shrink(arg))
  def pad(self, arg) -> Tensor: return Tensor(_pad(self.uop, tuple(arg)))
  def flip(self, axis: int) -> Tensor: return Tensor(flip_axis(self.uop, axis % self.ndim))
  def __getitem__(self, idx) -> Tensor:
    """Ints and step-1 slices only."""
    idx = idx if isinstance(idx, tuple) else (idx,)
    arg, squeeze = [], []
    for d, i in enumerate(idx):
      n = self.shape[d]
      if isinstance(i, int):
        i = i % n
        arg.append((i, i + 1)); squeeze.append(d)
      else:
        b, e, step = i.indices(n)
        assert step == 1, "only step-1 slices"
        arg.append((b, max(b, e)))
    arg += [(0, s) for s in self.shape[len(arg):]]
    t = self.shrink(arg)
    return t.reshape(tuple(s for d, s in enumerate(t.shape) if d not in squeeze))
  @staticmethod
  def cat(ts: list[Tensor], dim: int = 0) -> Tensor:
    """Concatenate via pad + add: each piece is padded into place, then summed."""
    dim = dim % ts[0].ndim
    total = sum(t.shape[dim] for t in ts)
    out, off = None, 0
    for t in ts:
      p = [(0, 0)] * t.ndim
      p[dim] = (off, total - off - t.shape[dim])
      piece = t.pad(p)
      out = piece if out is None else out + piece
      off += t.shape[dim]
    return out
  def stack_with(self, *others) -> Tensor: return Tensor(UOp(Ops.STACK, (self.uop, *(o.uop for o in others))))

  # ---- reductions -----------------------------------------------------------
  def _reduce(self, op: Ops, axis=None, keepdim=False) -> Tensor:
    axes = list(range(self.ndim)) if axis is None else [a % self.ndim for a in (axis if isinstance(axis, (tuple, list)) else [axis])]
    keep = [i for i in range(self.ndim) if i not in axes]
    t = self.permute(keep + axes)
    r = Tensor(t.uop.reduce(op, len(axes)))
    if keepdim: r = r.reshape(tuple(1 if i in axes else s for i, s in enumerate(self.shape)))
    return r
  def sum(self, axis=None, keepdim=False): return self._reduce(Ops.ADD, axis, keepdim)
  def max(self, axis=None, keepdim=False): return self._reduce(Ops.MAX, axis, keepdim)
  def min(self, axis=None, keepdim=False): return -((-self).max(axis, keepdim))
  def mean(self, axis=None, keepdim=False):
    s = (self if self.dtype.is_float else self.float()).sum(axis, keepdim)
    return s * (s.numel() / self.numel())
  def var(self, axis=None, keepdim=False):
    return (self - self.mean(axis, keepdim=True)).square().mean(axis, keepdim)
  def softmax(self, axis=-1):
    e = (self - self.max(axis, keepdim=True)).exp()
    return e / e.sum(axis, keepdim=True)
  def log_softmax(self, axis=-1):
    m = self - self.max(axis, keepdim=True)
    return m - m.exp().sum(axis, keepdim=True).log()
  def logsumexp(self, axis=-1):
    m = self.max(axis, keepdim=True)
    return (self - m).exp().sum(axis, keepdim=True).log() + m
  def argmax(self, axis=-1):
    """Index of the first maximum, as i32."""
    axis = axis % self.ndim
    n = self.shape[axis]
    idx = Tensor.arange(n).reshape((n,) + (1,) * (self.ndim - axis - 1))
    is_max = self.eq(self.max(axis, keepdim=True))
    return is_max.where(n - idx, 0).max(axis) * -1 + n

  # ---- linear algebra -------------------------------------------------------
  def matmul(self, w: Tensor) -> Tensor:
    """(..., M, K) @ (..., K, N): broadcast to (..., M, N, K), multiply, reduce K."""
    x = self if self.ndim > 1 else self.reshape(1, -1)
    wt = w.reshape(-1, 1) if w.ndim == 1 else w
    wt = wt.transpose()                                     # (..., N, K)
    x = x.reshape(x.shape[:-1] + (1, x.shape[-1]))          # (..., M, 1, K)
    wt = wt.reshape(wt.shape[:-2] + (1,) + wt.shape[-2:])   # (..., 1, N, K)
    out = (x * wt).sum(-1)
    if self.ndim == 1: out = out.reshape(out.shape[:-2] + out.shape[-1:])
    if w.ndim == 1: out = out.reshape(out.shape[:-1])
    return out
  __matmul__ = matmul
  def linear(self, w: Tensor, b: Tensor | None = None) -> Tensor:
    out = self @ w
    return out + b if b is not None else out

  def conv2d(self, w: Tensor, b: Tensor | None = None, padding: int = 0, stride: int = 1) -> Tensor:
    """x (N, C, H, W), w (O, C, KH, KW). im2col from shrinks, then a matmul."""
    x = self.pad(((0, 0), (0, 0), (padding, padding), (padding, padding))) if padding else self
    N, C, H, W = x.shape
    O, _, KH, KW = w.shape
    OH, OW = (H - KH) // stride + 1, (W - KW) // stride + 1
    cols = []
    for ky in range(KH):
      for kx in range(KW):
        win = x.shrink(((0, N), (0, C), (ky, ky + stride * (OH - 1) + 1), (kx, kx + stride * (OW - 1) + 1)))
        if stride > 1: win = win._stride2d(stride, OH, OW)
        cols.append(win)                                     # (N, C, OH, OW) each
    patches = cols[0].stack_with(*cols[1:]) if len(cols) > 1 else cols[0].unsqueeze(0)   # (KH*KW, N, C, OH, OW)
    patches = patches.permute(1, 3, 4, 2, 0).reshape(N, OH, OW, C * KH * KW)
    out = patches @ w.reshape(O, C * KH * KW).T              # (N, OH, OW, O)
    out = out.permute(0, 3, 1, 2)
    return out + b.reshape(1, O, 1, 1) if b is not None else out

  def _stride2d(self, s: int, OH: int, OW: int) -> Tensor:
    N, C, H, W = self.shape
    t = self.pad(((0, 0), (0, 0), (0, OH * s - H), (0, OW * s - W)))
    return t.reshape(N, C, OH, s, OW, s)[:, :, :, 0, :, 0]

  def max_pool2d(self, k: int = 2) -> Tensor:
    N, C, H, W = self.shape
    return self[:, :, :H // k * k, :W // k * k].reshape(N, C, H // k, k, W // k, k).max((3, 5))

  # ---- losses ---------------------------------------------------------------
  def one_hot(self, n: int) -> Tensor:
    return self.unsqueeze(-1).eq(Tensor.arange(n)).float()
  def cross_entropy(self, labels: Tensor) -> Tensor:
    """self: (B, C) logits, labels: (B,) i32."""
    return -(self.log_softmax(-1) * labels.one_hot(self.shape[-1])).sum(-1).mean()

  # ---- autodiff ------------------------------------------------------------
  def backward(self, params: list[Tensor]) -> Tensor:
    """Set p.grad = d(self)/d(p) for each p. Gradients are lazy Tensors like everything else."""
    assert self.shape == (), "backward needs a scalar"
    root = getattr(self, "_graph", self.uop)
    for p, g in zip(params, gradient(root, [p.uop for p in params])):
      p.grad = Tensor(g) if p.grad is None else p.grad + Tensor(g)
    return self

  # ---- staging -------------------------------------------------------------
  def contiguous(self) -> Tensor:
    """Ask for a materialization point: the kernelizer will write this into a temp buffer."""
    return Tensor(UOp(Ops.STAGE, (self.uop,)))


if os.environ.get("OPT", "1") != "0":
  from uopc.opt import HandOpt, SearchOpt
  Tensor.opt = SearchOpt() if os.environ.get("OPT") == "2" else HandOpt()
uopc/nn.py · 100 lines
uopc/nn.py
"""Week 5: layers and optimizers, written against the Tensor API only."""
from __future__ import annotations
import math
from uopc.tensor import Tensor


class Linear:
  def __init__(self, n_in: int, n_out: int, bias: bool = True):
    bound = 1 / math.sqrt(n_in)
    self.weight = Tensor.uniform(n_in, n_out, low=-bound, high=bound, requires_grad=True)
    self.bias = Tensor.zeros(n_out).realize() if bias else None
    if self.bias is not None: self.bias.requires_grad = True
  def __call__(self, x: Tensor) -> Tensor: return x.linear(self.weight, self.bias)


class Conv2d:
  def __init__(self, c_in: int, c_out: int, k: int, padding: int = 0, stride: int = 1):
    bound = 1 / math.sqrt(c_in * k * k)
    self.weight = Tensor.uniform(c_out, c_in, k, k, low=-bound, high=bound, requires_grad=True)
    self.bias = Tensor.zeros(c_out).realize()
    self.bias.requires_grad = True
    self.padding, self.stride = padding, stride
  def __call__(self, x: Tensor) -> Tensor: return x.conv2d(self.weight, self.bias, self.padding, self.stride)


class LayerNorm:
  def __init__(self, n: int, eps: float = 1e-5):
    self.weight, self.bias, self.eps = Tensor.ones(n).realize(), Tensor.zeros(n).realize(), eps
    self.weight.requires_grad = self.bias.requires_grad = True
  def __call__(self, x: Tensor) -> Tensor:
    mu = x.mean(-1, keepdim=True)
    xc = x - mu
    return xc * (xc.square().mean(-1, keepdim=True) + self.eps).rsqrt() * self.weight + self.bias


class Embedding:
  """A one-hot matmul: differentiable without a scatter op."""
  def __init__(self, vocab: int, dim: int):
    self.vocab, self.weight = vocab, Tensor.randn(vocab, dim) * 0.02
    self.weight = self.weight.realize(); self.weight.requires_grad = True
  def __call__(self, idx: Tensor) -> Tensor: return idx.one_hot(self.vocab) @ self.weight


def parameters(obj) -> list[Tensor]:
  """Every Tensor with requires_grad reachable through attributes, lists and dicts."""
  out, seen = [], set()
  def walk(o):
    if id(o) in seen: return
    seen.add(id(o))
    if isinstance(o, Tensor):
      if o.requires_grad: out.append(o)
    elif isinstance(o, (list, tuple)): [walk(x) for x in o]
    elif isinstance(o, dict): [walk(x) for x in o.values()]
    elif hasattr(o, "__dict__"): [walk(x) for x in vars(o).values()]
  walk(obj)
  return out


class Optimizer:
  def __init__(self, params: list[Tensor]): self.params = params
  def zero_grad(self):
    for p in self.params: p.grad = None


class SGD(Optimizer):
  def __init__(self, params, lr=0.01, momentum=0.0):
    super().__init__(params)
    self.lr, self.momentum = lr, momentum
    self.v = [Tensor.zeros(*p.shape).realize() for p in params] if momentum else []
  def step(self, *extra: Tensor):
    """Update every parameter. Extra tensors (e.g. the loss) are realized in the same batch."""
    news = []
    for i, p in enumerate(self.params):
      g = p.grad
      if self.momentum:
        self.v[i].assign(self.v[i] * self.momentum + g); g = self.v[i]
      news.append(p.assign(p.detach() - g * self.lr))
    Tensor.realize_all(*news, *self.v, *extra)


class AdamW(Optimizer):
  def __init__(self, params, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8, wd=0.0):
    super().__init__(params)
    self.lr, self.b1, self.b2, self.eps, self.wd, self.t = lr, b1, b2, eps, wd, 0
    self.m = [Tensor.zeros(*p.shape).realize() for p in params]
    self.v = [Tensor.zeros(*p.shape).realize() for p in params]
  def step(self, *extra: Tensor):
    self.t += 1
    # bias corrections are python floats baked into the graph as constants: each step's
    # graph differs only in two CONSTs, so we pass them as tensors to keep the JIT cache warm
    c1 = Tensor([self.lr / (1 - self.b1 ** self.t)])
    c2 = Tensor([1 / (1 - self.b2 ** self.t)])
    news = []
    for p, m, v in zip(self.params, self.m, self.v):
      g = p.grad
      m.assign(m * self.b1 + g * (1 - self.b1))
      v.assign(v * self.b2 + g.square() * (1 - self.b2))
      upd = m * c1.reshape(*([1] * p.ndim)) / ((v * c2.reshape(*([1] * p.ndim))).sqrt() + self.eps)
      news.append(p.assign(p.detach() * (1 - self.lr * self.wd) - upd))
    Tensor.realize_all(*news, *self.m, *self.v, *extra)
examples/mnist_mlp.py · 43 lines
examples/mnist_mlp.py
"""Week 5 milestone: train an MLP on MNIST with our own compiler.

  python examples/mnist_mlp.py            # ~94% test accuracy in 600 steps (~15 ms/step)
  DEBUG=2 python examples/mnist_mlp.py    # print every kernel with its time
"""
import sys, time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent))
import numpy as np
from datasets import mnist
from uopc.tensor import Tensor
from uopc.nn import Linear, AdamW, parameters


class MLP:
  def __init__(self):
    self.l1, self.l2 = Linear(784, 128), Linear(128, 10)
  def __call__(self, x: Tensor) -> Tensor: return self.l2(self.l1(x).relu())


def main(steps=600, bs=128):
  Xtr, Ytr, Xte, Yte = mnist()
  Tensor.manual_seed(0)
  model = MLP()
  opt = AdamW(parameters(model), lr=1e-3)
  rng = np.random.default_rng(0)
  t0 = time.perf_counter()
  for step in range(steps):
    idx = rng.integers(0, len(Xtr), bs)
    x, y = Tensor(Xtr[idx]), Tensor(Ytr[idx])
    loss = model(x).cross_entropy(y)
    opt.zero_grad()
    loss.backward(opt.params)
    opt.step(loss)
    if step % 100 == 0 or step == steps - 1:
      print(f"step {step:4d}  loss {loss.item():.4f}  {(time.perf_counter() - t0) / (step + 1) * 1e3:.1f} ms/step")
  pred = np.concatenate([model(Tensor(Xte[i:i + 1000])).argmax(-1).numpy() for i in range(0, len(Xte), 1000)])
  acc = (pred == Yte).mean()
  print(f"test accuracy {acc * 100:.2f}%")
  return acc


if __name__ == "__main__": main()
examples/datasets.py · 29 lines
examples/datasets.py
"""Tiny dataset loaders. Files are cached in ~/.cache/uopc/data."""
import gzip, urllib.request
from pathlib import Path
import numpy as np

DATA = Path.home() / ".cache" / "uopc" / "data"
MNIST_URL = "https://storage.googleapis.com/cvdf-datasets/mnist/"
SHAKESPEARE_URL = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"


def fetch(url: str, name: str) -> Path:
  DATA.mkdir(parents=True, exist_ok=True)
  p = DATA / name
  if not p.exists():
    print(f"downloading {url}")
    urllib.request.urlretrieve(url, p)
  return p


def mnist():
  """(X_train (60000, 784) f32 in [0,1], Y_train (60000,) i32, X_test, Y_test)"""
  def images(n): return np.frombuffer(gzip.open(fetch(MNIST_URL + n, n)).read(), np.uint8, offset=16).reshape(-1, 784).astype(np.float32) / 255
  def labels(n): return np.frombuffer(gzip.open(fetch(MNIST_URL + n, n)).read(), np.uint8, offset=8).astype(np.int32)
  return (images("train-images-idx3-ubyte.gz"), labels("train-labels-idx1-ubyte.gz"),
          images("t10k-images-idx3-ubyte.gz"), labels("t10k-labels-idx1-ubyte.gz"))


def shakespeare() -> str:
  return fetch(SHAKESPEARE_URL, "tinyshakespeare.txt").read_text()
Suggest an edit to this chapter