topos

An autodiff compiler stack, small enough to read.

Record a graph once. Then read it: run it, differentiate it, schedule it, emit it. The tape is the spec and nothing rewrites it. The weights are yours. Every faster reading has to match the slow, obvious one, bit for bit, and a claim is one assert away from proof.

cargo add topos

Start here Read one graph six ways Open the playground

0 Parameter[2, 2]0 Parameter [2, 2]1 Input[1, 2]1 Input [1, 2]2 MatMul[1, 2]2 MatMul [1, 2]3 Tanh[1, 2]3 Tanh [1, 2]4 Mul[1, 2]4 Mul [1, 2]5 Sum[]5 Sum []6 Leaf[]6 Leaf []7 Broadcast[1, 2] shape=[1, 2]7 Broadcast [1, 2] shape=[1, 2]8 Leaf[1, 2]8 Leaf [1, 2]9 Add[1, 2]9 Add [1, 2]10 Mul[1, 2]10 Mul [1, 2]11 Mul[1, 2]11 Mul [1, 2]12 Leaf[1, 2]12 Leaf [1, 2]13 Add[1, 2]13 Add [1, 2]14 Add[1, 2]14 Add [1, 2]15 Leaf[1, 2]15 Leaf [1, 2]16 Mul[1, 2]16 Mul [1, 2]17 Sub[1, 2]17 Sub [1, 2]18 Mul[1, 2]18 Mul [1, 2]19 Leaf[1, 2]19 Leaf [1, 2]20 Add[1, 2]20 Add [1, 2]21 Permute[2, 2] order=[1, 0]21 Permute [2, 2] order=[1, 0]22 MatMul[1, 2]22 MatMul [1, 2]23 Permute[2, 1] order=[1, 0]23 Permute [2, 1] order=[1, 0]24 MatMul[2, 2]24 MatMul [2, 2]25 Leaf[1, 2]25 Leaf [1, 2]26 Add[1, 2]26 Add [1, 2]27 Leaf[2, 2]27 Leaf [2, 2]28 Add[2, 2]28 Add [2, 2]
Plate 1. A matrix product, a tanh, and a loss, with the derivative appended: the lighter nodes are the chain rule, recorded as more of the same graph. Hover a node to follow its edges.

The tape is the spec

A typical compiler rewrites the program until it can run. This one does not. Recording writes a list of operations, one line each, with every shape inferred at the line that records it. Sealing freezes it. Everything after — gradients, schedules, fusions, backends, emitted text — is a named way of reading that same list. You can print each reading. None of it is magic.

The interpreter's bits are the truth

The plain interpreter is the executable spec. Its bits are the same in every build, on every platform. A compiled plan under Exact must reproduce them; a backend that cannot is declined, not trusted. Reordering float math is a labeled choice on the entry, never a silent effect of a feature flag. Two identical runs cannot differ, so a rerun is a checksum.

Built for learning and research

Learners and teachers first: the spec, its gradient, the schedule, and the emitted text are readings of one graph, all printable. Systems researchers second: a new element type, AD mode, fusion, backend, or emission target plugs in at a named seam and is graded against the interpreter, without forking the crate. The core stays closed and simple on purpose.

Six lines, then everything else is a reading#

The whole library fits in one loop. Record the graph in a closure; the return value is the set of names that leave the tape. Feed a sample, ask for the gradient of one scalar, step the parameters. The network never changes.

use topos::{Detach, Tape, Tensor};

let (network, [w, x, y, loss]) = Tape::record(|tape| {
    let w = tape.parameter(0.0_f64);
    let x = tape.input(0.0);
    let y = tape.input(0.0);
    let error = w * x - y;
    [w, x, y, error * error].detach()
});
let mut parameters = network.parameters();

let samples = [(1.0, 2.0), (2.0, 4.0), (3.0, 6.0)];
for step in 0..100 {
    let (sample_x, sample_y) = samples[step % samples.len()];
    let run = network.forward(&parameters, [(x, sample_x.into()), (y, sample_y.into())]);
    let gradients = run.backward(loss).parameters(&parameters);
    parameters = parameters.step(&gradients, |w, g| w.clone() - g.clone() * Tensor::from(0.02));
}
assert!((parameters.of(w).scalar() - 2.0).abs() < 1e-6);
   0  Parameter                         []
   1  Input                             []
   2  Input                             []
   3  Mul            0, 1               []
   4  Sub            3, 2               []
   5  Mul            4, 4               []
tape: 6 nodes, 1 parameter, 2 inputs
Plate 2. What that closure recorded, printed by describe: the spec. Hover a line.
a hundred steps toward y = 2x ── w = 2.00000 s 1 ┤ q │ u │ a 10⁻⁴ ┤ r │ e 10⁻⁸ ┤ d │ 10⁻¹² ┤ e │ r 10⁻¹⁶ ┤ r │ └┬───────────┬────────────┬───────────┬────────────┬───────────┬ 0 20 40 60 80 100 step
Plate 3. What the loop did to the one table it holds. The graph is the same six lines before and after.

One spec, a short list of interpretations#

This table is the compiler. It is small enough to read because there is nothing else, and every row is on this site with a plate.

readingwhat answers itand prints
specTape, Networkone line per node — recording
valueNetwork::forward, BoundEntry::interpretthe exact oracle over the whole spec, or over a declared closure
cotangentRun::backwardthe engine reverse scan — differentiation
traceTape::differentiatethe same rules recording themselves as more spec
scheduleBoundEntry::lowera Plan: keep-set, liveness, elections — entries and plans
catalogPlan::patternsthe offers a run elected, as data, never as rewrites
textPlan::emit_stablehlothe plan for an industrial compiler — emission

What the readings look like#

   0  Parameter                         [2, 2]
   1  Input                             [1, 2]
   2  MatMul         1, 0               [1, 2]
   3  Tanh           2                  [1, 2]
   4  Mul            3, 3               [1, 2]
   5  Sum            4                  []
23 nodes appended by differentiate; node 28 is d loss / d weights
   6  Leaf                              []
   7  Broadcast      6  shape=[1, 2]    [1, 2]
   8  Leaf                              [1, 2]
   9  Add            8, 7               [1, 2]
  10  Mul            9, 3               [1, 2]
  11  Mul            9, 3               [1, 2]
  12  Leaf                              [1, 2]
  13  Add            12, 10             [1, 2]
  14  Add            13, 11             [1, 2]
  15  Leaf                              [1, 2]
  16  Mul            3, 3               [1, 2]
  17  Sub            15, 16             [1, 2]
  18  Mul            14, 17             [1, 2]
  19  Leaf                              [1, 2]
  20  Add            19, 18             [1, 2]
  21  Permute        0  order=[1, 0]    [2, 2]
  22  MatMul         20, 21             [1, 2]
  23  Permute        1  order=[1, 0]    [2, 1]
  24  MatMul         23, 20             [2, 2]
  25  Leaf                              [1, 2]
  26  Add            25, 22             [1, 2]
  27  Leaf                              [2, 2]
  28  Add            27, 24             [2, 2]
tape: 29 nodes, 1 parameter, 1 input
Plate 4. differentiate appended twenty-three nodes to the same tape. Nothing above the divider changed; the last line is d loss / d weights.

The gradient is more spec

Derivative knowledge lives in exactly one rule body. One reading of it computes; the other records the same steps as ordinary nodes, below a divider. The gradient graph is printable, plannable, emittable, and — being ordinary differentiable nodes — differentiable again. Second derivatives are the absence of a wall.

   0  Parameter                         [2, 2]       freed after 2
   1  Input                             [1, 2]       freed after 2
   2  MatMul         1, 0               [1, 2]       freed after 3
   3  Tanh           2                  [1, 2]       freed after 4
   4  Mul            3, 3               [1, 2]       freed after 5
   5  Sum            4                  []           kept
plan: forward; 6 of 6 nodes evaluated, 1 readable
live volume: peak 8 elements at node 2, retain-all 13
Plate 5. Observability is declared. Lowering the same graph with hidden observed turns one line from freed to kept; nothing else changes.

Say what you want to read

An entry names its roots and the interiors it will also read. That declaration is the license every optimization runs on: the declared must survive, the rest may be freed, skipped, or fused. Add one name and a line turns from freed to kept. Nothing else moves.

module @topos {
  func.func @main(%arg0: tensor<2x1x3x3xf32>, %arg1: tensor<2xf32>, %arg2: tensor<1x1x6x6xf32>) -> (tensor<1x2x2x2xf32>) {
    %v7 = stablehlo.transpose %arg0, dims = [1, 2, 3, 0] : (tensor<2x1x3x3xf32>) -> tensor<1x3x3x2xf32>
    %v8 = stablehlo.reshape %v7 : (tensor<1x3x3x2xf32>) -> tensor<9x2xf32>
    %v9_kernel = stablehlo.reshape %v8 : (tensor<9x2xf32>) -> tensor<1x3x3x2xf32>
    %v9_windows = stablehlo.convolution(%arg2, %v9_kernel) dim_numbers = [b, f, 0, 1]x[i, 0, 1, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[0, 0], [0, 0]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x1x6x6xf32>, tensor<1x3x3x2xf32>) -> tensor<1x4x4x2xf32>
    %v9 = stablehlo.reshape %v9_windows : (tensor<1x4x4x2xf32>) -> tensor<16x2xf32>
    %v10 = stablehlo.broadcast_in_dim %arg1, dims = [1] : (tensor<2xf32>) -> tensor<16x2xf32>
    %v11 = stablehlo.add %v9, %v10 : tensor<16x2xf32>
    %v12 = stablehlo.reshape %v11 : (tensor<16x2xf32>) -> tensor<1x4x4x2xf32>
    %v13 = stablehlo.transpose %v12, dims = [0, 3, 1, 2] : (tensor<1x4x4x2xf32>) -> tensor<1x2x4x4xf32>
    %v14 = stablehlo.constant dense<0.0> : tensor<1x2x4x4xf32>
    %v15 = stablehlo.maximum %v13, %v14 : tensor<1x2x4x4xf32>
    %v27_seed = stablehlo.constant dense<0xFF800000> : tensor<f32>
    %v27 = "stablehlo.reduce_window"(%v15, %v27_seed) ({
    ^bb0(%v27_left: tensor<f32>, %v27_right: tensor<f32>):
      %v27_larger = stablehlo.maximum %v27_left, %v27_right : tensor<f32>
      stablehlo.return %v27_larger : tensor<f32>
    }) {window_dimensions = array<i64: 1, 1, 2, 2>, window_strides = array<i64: 1, 1, 2, 2>} : (tensor<1x2x4x4xf32>, tensor<f32>) -> tensor<1x2x2x2xf32>
    return %v27 : tensor<1x2x2x2xf32>
  }
}
Plate 6. A convolution recorded as pad, unfold, permute, reshape, and a matrix product, emitted as stablehlo.convolution; a max-pool as reduce_window. The same match that fuses at home raises abroad.

Fusion raises

A convolution is recorded as pad, unfold, permute, reshape, and one matrix product. The catalog finds that idiom again and, at home, calls a fused kernel; abroad, it writes stablehlo.convolution. Speed lives in the catalog, not in a new opcode with its own backward rule.

the learned decision surface •• class +1 •• class -1 1.5 ┤ │ 2 1.0 ┤ │ ▓ 1 0.5 ┤ ▓ │ ▓ │ ▒ 0 0.0 ┤ ▒ │ ▒ -1 -0.5 ┤ ░ │ ░ -2 -1.0 ┤ ░ └┬──────┬───────┬──────┬───────┬───────┬──────┬───────┬──────┬ -1.5 -1.0 -0.5 0.0 0.5 1.0 1.5 2.0 2.5
Plate 7. The two moons, separated: the decision surface of a [2, 16, 16, 1] MLP after two thousand full-batch steps, drawn by a grid-shaped twin of the training expression.

The state is yours

Training is a pure data transform of a caller-owned table. A what-if costs one clone. A decision surface is a grid-shaped twin of the training expression, recorded once beside it and fed the raster. The gallery runs this from a bigram table to GPT-2.

the spec over f64, f32, and Bf16 is the same text: true
loss over f64   3.7912668765691295
loss over f32   3.791267
loss over Bf16  3.78125
Plate 8. One recording over three elements. The spec is the same text over each; the run's bits are each number's own.

The element is the seam

The graph is always tensors; a scalar is rank 0. The open plug is a number: implement two traits on your own float — a dual, an interval, a research format — and the tensor machinery, the derivative rules, the plans, and the notebook come along. Bf16 is built in. Forward mode was written outside the crate this way.

plan
   0  Parameter                         [8, 16]      freed after 5
   1  Parameter                         [16]         freed after 6
   2  Parameter                         [16, 4]      freed after 24
   3  Parameter                         [4]          freed after 10
   4  Input                             [32, 8]      freed after 47
   5  MatMul         4, 0               [32, 16]     freed after 7
   6  BroadcastAlong 1  axis=0 extent=32 [32, 16]    freed after 7
   7  Add            5, 6               [32, 16]     freed after 8
   8  Tanh           7                  [32, 16]     freed after 33
   9  MatMul         8, 2               [32, 4]      freed after 11
  10  BroadcastAlong 3  axis=0 extent=32 [32, 4]     freed after 11
  11  Add            9, 10              [32, 4]      freed after 12
  12  Sum            11                 []           kept
  13  Leaf                              []           freed after 14
  14  Broadcast      13  shape=[32, 4]  [32, 4]      freed after 16
  15  Leaf                              [32, 4]      freed after 16
  16  Add            15, 14             [32, 4]      freed after 20
  17  Leaf                              [32, 4]      freed after 18
  18  Add            17, 16             [32, 4]      freed after 27
  19  Leaf                              [32, 4]      freed after 20
  20  Add            19, 16             [32, 4]      freed after 21
  21  SumAlong       20  axis=0         [4]          freed after 23
  22  Leaf                              [4]          freed after 23
  23  Add            22, 21             [4]          kept
  24  Permute        2  order=[1, 0]    [4, 16]      freed after 25
  25  MatMul         18, 24             [32, 16]     freed after 29
  26  Permute        8  order=[1, 0]    [16, 32]     freed after 27
  27  MatMul         26, 18             [16, 4]      freed after 31
  28  Leaf                              [32, 16]     freed after 29
  29  Add            28, 25             [32, 16]     freed after 35
  30  Leaf                              [16, 4]      freed after 31
  31  Add            30, 27             [16, 4]      kept
  32  Leaf                              [32, 16]     freed after 34
  33  Mul            8, 8               [32, 16]     freed after 34
  34  Sub            32, 33             [32, 16]     freed after 35
  35  Mul            29, 34             [32, 16]     freed after 37
  36  Leaf                              [32, 16]     freed after 37
  37  Add            36, 35             [32, 16]     freed after 41
  38  Leaf                              [32, 16]     freed after 39
  39  Add            38, 37             [32, 16]     freed after 48
  40  Leaf                              [32, 16]     freed after 41
  41  Add            40, 37             [32, 16]     freed after 42
  42  SumAlong       41  axis=0         [16]         freed after 44
  43  Leaf                              [16]         freed after 44
  44  Add            43, 42             [16]         kept
  47  Permute        4  order=[1, 0]    [8, 32]      freed after 48
  48  MatMul         47, 39             [8, 16]      freed after 52
  51  Leaf                              [8, 16]      freed after 52
  52  Add            51, 48             [8, 16]      kept
plan: forward; 49 of 53 nodes evaluated, 5 readable
live volume: peak 2373 elements at node 29, retain-all 11922

                         live volume (elements)
                                ── live
2500 ┤                                      ⡄    ⡤⢤        ⢠
     │                                     ⡸⢱   ⢰⠁ ⢇       ⡎⡆
2000 ┤         ⡀                        ⡤⠒⢤⠃⠸⡀  ⡎  ⠘⡄ ⢀  ⣀⣰⠁⢇
     │        ⢠⢳                       ⢰⠁    ⡇ ⡰⠁   ⢣ ⡎⡆⢠⠃  ⢸
1500 ┤        ⡜ ⢇                      ⡎     ⢱⠴⠁    ⠘⡼ ⠸⡜   ⠈⡆
     │       ⢠⠃ ⠈⢆ ⢀⡠⡀    ⣀⢄⡀⡠⠤⠔⠑⡄    ⡜              ⠁  ⠁    ⢣
1000 ┤      ⢠⠃   ⠈⠖⠁ ⠱⡀ ⡠⠊  ⠈    ⠈⠒⠤⠤⠚                        ⢇ ⡰⠑⢢
     │     ⢀⠇         ⠈⠊                                      ⠈⠉  ⠈⡆
 500 ┤     ⡜                                                       ⢇ ⡀
     │  ⢀⣀⡰⠁                                                       ⠘⠉
   0 ┤⠒⠒⠁
     └┬────────────┬────────────┬────────────┬────────────┬────────────┬
      0           10           20           30           40           50
                                scheduled node
Plate 9. The plan card.

It draws itself

In an Evcxr notebook, the last expression of a cell is a card: a spec dump, a schedule with its live volume, a tensor as a shaded table, a backward pass as a curve of norms. The cards are ordinary to_html strings, and this page embeds them as emitted.

Where to go#

Start here

  • Getting startedInstall, the first tape, a gradient, a training step, and the readings of one small graph.
  • The stack, read six waysThe flagship chapter: one graph as spec, derivative, plan, StableHLO, two runs that must agree, and a walk by hand.
  • PlaygroundWrite a small graph in the browser; every reading remakes itself as you type.

The guide

  • RecordingTapes, values, symbols, shapes at record time, sealing and reopening.
  • DifferentiationEngine scan, recorded reverse mode, seeds, second derivatives, forward mode from outside.
  • TrainingSpec and state, learning rates as clones, optimizers, twins, checkpoints.
  • Entries and plansObservability, liveness, skipped nodes, patterns, Exact and Fast.
  • EmissionThe plan as StableHLO, raises, result order, conformance.
  • The neural tierEvery layer is a spelling over the public operations, and the dump proves it.
  • The element seamBring your own number.
  • AccelerationOpt-in backends, the measured ladder, and the tally.
  • NotebooksEvcxr, two rules, and the cards.

Examples

  • The galleryEvery example with what it printed and its source: a scalar chain, moons, makemore act by act, MNIST, CIFAR-10, GPT-2, Llama.
  • GPT-2 on toposFour engines, one text, and a GPU plugin caught being wrong.
  • Llama on toposTinyLlama and Llama 2 7B from one module tree.

Why it is shaped this way

Reference

  • TerminologyThe vocabulary contract, illustrated.
  • API referenceTwo maps: topos::model to train, topos::compiler to inspect and emit.
  • ChangelogEvery release, written for a person.

The name#

A topos is a place. Here it is the one where the whole compiler stack stays in view: the spec, its gradient, the schedule, and the emitted text are readings of one graph, and every claim is one assert from proof.

live volume along the schedule ── live elements ── peak 15 e 16 ┤ l │ e │ ╭───╮ ╭──╮ m 12 ┤ ╭╮ ╭───╯ ╰─╮ ╭─╯ ╰─╮ ╭╮ ╭─── e │ ╭─╯╰───╯ ╰──╯ ╰───╯╰───╯ n 8 ┤ ╭╮ ╭──╯ t │ ╭─╯╰─╮ ╭───╮╭─╯ s │ ╭─╯ ╰──╯ ╰╯ 4 ┤ ╯ └┬────────────┬────────────┬────────────┬────────────┬────────────┬ 0 5 10 15 20 25 scheduled node
Plate 10. The plan's own memory accounting, drawn from Plan::live_series: elements alive after each scheduled node of the walkthrough's training step.