topos
Start here
The stack, read six ways
One graph — a matrix product, a tanh, a loss — as the spec, its derivative, a plan, StableHLO, two runs that must agree, and a walk by hand.
Topos is an autodiff compiler, and both halves of that phrase show up on the smallest graph that still has something to say: a matrix product, a tanh, and a loss. The loss is one number that says how wrong the output is; training pushes it down, and the gradient — how much the loss moves when each weight moves — is what tells training which way to push.
Recording writes the graph down once, as a list of operations. That list is the tape, and once sealed it is the spec: the one definition of what the program means. Everything after that is a different way of reading the same list, and this page prints each reading in turn. Nothing is rewritten between them; the spec never changes. The whole chapter is also one program, examples/walkthrough.rs: cargo run --example walkthrough prints every plate below.
use topos::{Tape, Tensor};
let tape: Tape<f32> = Tape::new();
let weights = tape.parameter(Tensor::new([2, 2], [0.5_f32, -0.25, 0.25, 0.4]));
let input = tape.input(Tensor::new([1, 2], [1.0_f32, 2.0]));
let hidden = input.matmul(weights).tanh();
let loss = (hidden * hidden).sum();weights is a parameter, the thing training will change; input is fed per run and carries a default, which this page uses throughout.
1. The spec#
One line per recorded operation: its number, its opcode, the numbers of the operands it reads, and its shape. Shapes were checked as each line was written; a [1, 2] by [3, 3] product would have panicked at the line that tried to record it, with both shapes in the message.
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 [] tape: 6 nodes, 1 parameter, 1 input
The code that produced it
use topos::{Tape, Tensor};
let tape: Tape<f32> = Tape::new();
let weights = tape.parameter(Tensor::new([2, 2], [0.5_f32, -0.25, 0.25, 0.4]));
let input = tape.input(Tensor::new([1, 2], [1.0_f32, 2.0]));
let hidden = input.matmul(weights).tanh();
let _loss = (hidden * hidden).sum();
show::spec(tape.describe())hidden read twice by the square.2. The derivative, as more spec#
A Value borrows the tape and cannot outlive it. A Symbol is the detached name every later phase speaks, and differentiate takes symbols (or values, which convert): the target, and the names to differentiate with respect to.
let adjoints = tape.differentiate(loss, [weights]);The chain rule is applied by the same rules the interpreter's reverse scan uses; here they record instead of compute. The gradient becomes ordinary nodes: printable, schedulable, emittable, and differentiable again.
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
differentiate appended twenty-three nodes to the same tape. Nothing above the divider changed; the last line is d loss / d weights.The code that produced it
use topos::{Tape, Tensor};
let tape: Tape<f32> = Tape::new();
let weights = tape.parameter(Tensor::new([2, 2], [0.5_f32, -0.25, 0.25, 0.4]));
let input = tape.input(Tensor::new([1, 2], [1.0_f32, 2.0]));
let hidden = input.matmul(weights).tanh();
let loss = (hidden * hidden).sum();
let spec_nodes = tape.len();
let adjoints = tape.differentiate(loss, [weights]);
let appended = tape.len() - spec_nodes;
let gradient = adjoints.of(weights.symbol()).index();
show::spec_from(
tape.describe(),
spec_nodes,
&format!("{appended} nodes appended by differentiate; node {gradient} is d loss / d weights"),
)How to read the appended lines:
- a Leaf holding one is the seed, spread by
Broadcastover the shape theSumreduced; hidden * hiddengets the product rule once per operand — the twoMul 9, 3lines;1 - hidden * hiddenis the derivative of tanh — theLeaf,Mul 3, 3,Sub;- the matmul rule pushes the result back as two transposed products, one per operand:
PermutethenMatMul, twice; - each Leaf-then-Add pair starts a gradient at zero and adds one consumer's share. That is how a value used twice collects both contributions, and it is stated once, in the engine, not in each rule.
Sealing turns the tape into an immutable Network. From here on nothing writes the graph: training only replaces the caller-owned Parameters, which start from the recorded initial values.
let (hidden, loss) = (hidden.symbol(), loss.symbol());
let network = tape.into_network();
let parameters = network.parameters();3. The plan#
An entry declares what a run may read: its roots — the loss and the gradient, in that order, which is what adjoints.roots() lists — plus any interior we also want to look at. Reading an undeclared interior panics. That declaration is the license the scheduler works under: the declared must survive; the undeclared may be skipped, freed, or fused.
let plan = network.entry(adjoints.roots()).observe([hidden]).lower();
println!("{}", plan.describe());Each line of the plan also says what happens to the value once computed: kept for reading, or freed after its last use. The last lines are the plan's own memory accounting: the most elements alive at once, against the total if nothing were ever freed.
0 Parameter [2, 2] freed after 2 1 Input [1, 2] freed after 23 2 MatMul 1, 0 [1, 2] freed after 3 3 Tanh 2 [1, 2] kept 4 Mul 3, 3 [1, 2] freed after 5 5 Sum 4 [] kept 6 Leaf [] freed after 7 7 Broadcast 6 shape=[1, 2] [1, 2] freed after 9 8 Leaf [1, 2] freed after 9 9 Add 8, 7 [1, 2] freed after 11 10 Mul 9, 3 [1, 2] freed after 13 11 Mul 9, 3 [1, 2] freed after 14 12 Leaf [1, 2] freed after 13 13 Add 12, 10 [1, 2] freed after 14 14 Add 13, 11 [1, 2] freed after 18 15 Leaf [1, 2] freed after 17 16 Mul 3, 3 [1, 2] freed after 17 17 Sub 15, 16 [1, 2] freed after 18 18 Mul 14, 17 [1, 2] freed after 20 19 Leaf [1, 2] freed after 20 20 Add 19, 18 [1, 2] freed after 24 23 Permute 1 order=[1, 0] [2, 1] freed after 24 24 MatMul 23, 20 [2, 2] freed after 28 27 Leaf [2, 2] freed after 28 28 Add 27, 24 [2, 2] kept plan: forward; 25 of 29 nodes evaluated, 3 readable live volume: peak 15 elements at node 28, retain-all 56
The code that produced it
use topos::{Tape, Tensor};
let tape: Tape<f32> = Tape::new();
let weights = tape.parameter(Tensor::new([2, 2], [0.5_f32, -0.25, 0.25, 0.4]));
let input = tape.input(Tensor::new([1, 2], [1.0_f32, 2.0]));
let hidden = input.matmul(weights).tanh();
let loss = (hidden * hidden).sum();
let adjoints = tape.differentiate(loss, [weights]);
let hidden = hidden.symbol();
let network = tape.into_network();
// The entry declares what a run may read: the loss and the
// gradient as roots, plus `hidden` as an observed interior.
let plan = network.entry(adjoints.roots()).observe([hidden]).lower();
let results = plan.results().iter().map(|symbol| symbol.index());
show::spec_marked(plan.describe(), results)Four nodes are missing: 21, 22, 25, 26. The rules also produced the gradient with respect to the input; nobody declared it, so it is never computed. Fusion is an offer the catalog makes for shapes it recognizes — windows, batch-norm — and a dense layer this small runs as primitives, so no patterns were elected here. The plans page has a convolution that does elect two.
The code that produced it
use malevich::{Line, LineStyle, Plot, Rule};
use topos::{Tape, Tensor};
let tape: Tape<f32> = Tape::new();
let weights = tape.parameter(Tensor::new([2, 2], [0.5_f32, -0.25, 0.25, 0.4]));
let input = tape.input(Tensor::new([1, 2], [1.0_f32, 2.0]));
let hidden = input.matmul(weights).tanh();
let loss = (hidden * hidden).sum();
let adjoints = tape.differentiate(loss, [weights]);
let hidden = hidden.symbol();
let network = tape.into_network();
let plan = network.entry(adjoints.roots()).observe([hidden]).lower();
let live: Vec<f64> = plan.live_series().iter().map(|&elements| elements as f64).collect();
let peak = live.iter().copied().fold(0.0, f64::max);
show::plot(
Plot::new()
.layer(Line::y(live).label("live elements").style(LineStyle::Corners))
.layer(Rule::h(peak).label(format!("peak {peak}")))
.title("live volume along the schedule")
.x_label("scheduled node")
.y_label("elements"),
)4. The plan as StableHLO#
Emission is describe in another syntax: the same schedule, written as text an industrial compiler accepts. Nothing XLA-shaped lives in the crate; the text is the boundary.
println!("{}", plan.emit_stablehlo().expect("every operation in this graph lowers"));module @topos {
func.func @main(%arg0: tensor<2x2xf32>, %arg1: tensor<1x2xf32>) -> (tensor<f32>, tensor<2x2xf32>, tensor<1x2xf32>) {
%v2 = stablehlo.dot_general %arg1, %arg0, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x2xf32>) -> tensor<1x2xf32>
%v3 = stablehlo.tanh %v2 : tensor<1x2xf32>
%v4 = stablehlo.multiply %v3, %v3 : tensor<1x2xf32>
%v5_seed = stablehlo.constant dense<0.0> : tensor<f32>
%v5 = stablehlo.reduce(%v4 init: %v5_seed) applies stablehlo.add across dimensions = [0, 1] : (tensor<1x2xf32>, tensor<f32>) -> tensor<f32>
%v6 = stablehlo.constant dense<1.0> : tensor<f32>
%v7 = stablehlo.broadcast_in_dim %v6, dims = [] : (tensor<f32>) -> tensor<1x2xf32>
%v8 = stablehlo.constant dense<0.0> : tensor<1x2xf32>
%v9 = stablehlo.add %v8, %v7 : tensor<1x2xf32>
%v10 = stablehlo.multiply %v9, %v3 : tensor<1x2xf32>
%v11 = stablehlo.multiply %v9, %v3 : tensor<1x2xf32>
%v12 = stablehlo.constant dense<0.0> : tensor<1x2xf32>
%v13 = stablehlo.add %v12, %v10 : tensor<1x2xf32>
%v14 = stablehlo.add %v13, %v11 : tensor<1x2xf32>
%v15 = stablehlo.constant dense<1.0> : tensor<1x2xf32>
%v16 = stablehlo.multiply %v3, %v3 : tensor<1x2xf32>
%v17 = stablehlo.subtract %v15, %v16 : tensor<1x2xf32>
%v18 = stablehlo.multiply %v14, %v17 : tensor<1x2xf32>
%v19 = stablehlo.constant dense<0.0> : tensor<1x2xf32>
%v20 = stablehlo.add %v19, %v18 : tensor<1x2xf32>
%v23 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<1x2xf32>) -> tensor<2x1xf32>
%v24 = stablehlo.dot_general %v23, %v20, contracting_dims = [1] x [0] : (tensor<2x1xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>
%v27 = stablehlo.constant dense<0.0> : tensor<2x2xf32>
%v28 = stablehlo.add %v27, %v24 : tensor<2x2xf32>
return %v5, %v28, %v3 : tensor<f32>, tensor<2x2xf32>, tensor<1x2xf32>
}
}The code that produced it
use topos::{Tape, Tensor};
let tape: Tape<f32> = Tape::new();
let weights = tape.parameter(Tensor::new([2, 2], [0.5_f32, -0.25, 0.25, 0.4]));
let input = tape.input(Tensor::new([1, 2], [1.0_f32, 2.0]));
let hidden = input.matmul(weights).tanh();
let loss = (hidden * hidden).sum();
let adjoints = tape.differentiate(loss, [weights]);
let hidden = hidden.symbol();
let network = tape.into_network();
let plan = network.entry(adjoints.roots()).observe([hidden]).lower();
show::hlo(plan.emit_stablehlo().expect("every operation in this graph lowers"))The function takes the parameter and the input, in recording order, and returns exactly what the entry declared, in that order: the loss, the gradient, then hidden. The four skipped nodes are not in the text either.
5. Two runs that must agree#
The interpreter over the whole spec is the oracle. Its bits are the same in every build, on every platform, and they define what the answer is. Its reverse scan applies the derivative rules directly, without recording.
let oracle = network.forward(¶meters, []);
let engine_gradient = oracle.backward(loss).parameters(¶meters);A plan under Exact must reproduce those bits. Under Fast, the default, a backend feature may serve large products with sums in a different order; that difference is a labeled choice on the entry, never a silent effect of a feature flag.
let declare = || network.entry(adjoints.roots()).observe([hidden]);
let exact_run = declare().numerics(Numerics::Exact).lower().forward(¶meters, []);
let fast_run = declare().lower().forward(¶meters, []);
let recorded_gradient = exact_run.recorded_gradients(&adjoints);loss: interpreter 0.83054614 exact plan 0.83054614 fast plan 0.83054614 d loss / d weights: engine scan [[0.6397, 0.75025934], [1.2794, 1.5005187]] recorded gradient [[0.6397, 0.75025934], [1.2794, 1.5005187]] hidden, declared readable: [[0.7615942, 0.5005202]] exact plan == interpreter, bit for bit: asserted recorded gradient == engine scan, bit for bit: asserted fast plan == exact plan: yes; no backend claimed any of this graph in this build
Exact plan reproduces them; the recorded gradient reproduces the engine scan. All of it asserted.The code that produced it
use topos::{Numerics, Tape, Tensor};
let tape: Tape<f32> = Tape::new();
let weights = tape.parameter(Tensor::new([2, 2], [0.5_f32, -0.25, 0.25, 0.4]));
let input = tape.input(Tensor::new([1, 2], [1.0_f32, 2.0]));
let hidden = input.matmul(weights).tanh();
let loss = (hidden * hidden).sum();
let adjoints = tape.differentiate(loss, [weights]);
let (weights, hidden, loss) = (weights.symbol(), hidden.symbol(), loss.symbol());
let network = tape.into_network();
let parameters = network.parameters();
// The interpreter over the whole spec is the oracle; its reverse
// scan applies the derivative rules without recording.
let oracle = network.forward(¶meters, []);
let engine_gradient = oracle.backward(loss).parameters(¶meters);
// A plan under `Exact` must reproduce the oracle's bits; the
// default `Fast` plan may let a backend reorder a sum.
let declare = || network.entry(adjoints.roots()).observe([hidden]);
let exact_run = declare().numerics(Numerics::Exact).lower().forward(¶meters, []);
let fast_run = declare().lower().forward(¶meters, []);
let recorded_gradient = exact_run.recorded_gradients(&adjoints);
let bits = |tensor: &Tensor<f32>| -> Vec<u32> { tensor.to_vec().iter().map(|v| v.to_bits()).collect() };
assert_eq!(bits(oracle.of(loss)), bits(exact_run.of(loss)));
assert_eq!(bits(engine_gradient.of(weights)), bits(recorded_gradient.of(weights)));
let fast_agrees = bits(fast_run.of(loss)) == bits(exact_run.of(loss));
show::lines([
"loss:".to_string(),
format!(" interpreter {}", oracle.of(loss)),
format!(" exact plan {}", exact_run.of(loss)),
format!(" fast plan {}", fast_run.of(loss)),
"d loss / d weights:".to_string(),
format!(" engine scan {}", engine_gradient.of(weights)),
format!(" recorded gradient {}", recorded_gradient.of(weights)),
format!("hidden, declared readable: {}", exact_run.of(hidden)),
String::new(),
"exact plan == interpreter, bit for bit: asserted".to_string(),
"recorded gradient == engine scan, bit for bit: asserted".to_string(),
format!(
"fast plan == exact plan: {}",
if fast_agrees { "yes; no backend claimed any of this graph in this build" } else { "no; a backend served part of it under Fast" }
),
])A few steps of gradient descent close the loop. The plan was lowered once and holds no state: each step feeds the current parameters, reads the recorded gradient, and mints the next parameters. The network never changes, and a clone of the state is all a what-if costs. Nothing broadcasts by accident, so the scalar learning rate is spread over each gradient explicitly.
The code that produced it
use malevich::{Line, Plot};
use topos::{Tape, Tensor};
let tape: Tape<f32> = Tape::new();
let weights = tape.parameter(Tensor::new([2, 2], [0.5_f32, -0.25, 0.25, 0.4]));
let input = tape.input(Tensor::new([1, 2], [1.0_f32, 2.0]));
let hidden = input.matmul(weights).tanh();
let loss = (hidden * hidden).sum();
let adjoints = tape.differentiate(loss, [weights]);
let loss = loss.symbol();
let network = tape.into_network();
// Lowered once, the plan holds no state: each step feeds the
// current parameters, reads the recorded gradient, and mints the
// next parameters. The learning rate is spread over each gradient
// explicitly, because nothing broadcasts by accident.
let plan = network.entry(adjoints.roots()).lower();
let learning_rate = Tensor::from(0.05_f32);
let mut current = network.parameters();
let mut losses = Vec::new();
for _ in 0..40 {
let run = plan.forward(¤t, []);
losses.push(run.of(loss).scalar());
let gradients = run.recorded_gradients(&adjoints);
current = current.step(&gradients, |weight, gradient| {
weight.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
});
}
show::plot(
Plot::new()
.layer(Line::y(losses).label("loss"))
.title("forty steps through the compiled plan")
.x_label("step")
.y_label("loss"),
)6. The spec is executable#
The printed lines are an instruction list, and walking them is the interpreter: sources take their stored payloads, computed nodes express their opcode over operands computed earlier. Opcode::express is public, so the walk can be written outside the crate — and it runs under the exact posture, as the whole-spec oracle does by construction.
walked 6 lines by hand, outside the crate loss by hand: 0.83054614 loss by interpreter: 0.83054614 same bits: asserted
Opcode::express, outside the crate, reproduces the interpreter's loss bit for bit.The code that produced it
use topos::{Numerics, Tape, Tensor};
let tape: Tape<f32> = Tape::new();
let weights = tape.parameter(Tensor::new([2, 2], [0.5_f32, -0.25, 0.25, 0.4]));
let input = tape.input(Tensor::new([1, 2], [1.0_f32, 2.0]));
let hidden = input.matmul(weights).tanh();
let loss = (hidden * hidden).sum().symbol();
let network = tape.into_network();
let oracle = network.forward(&network.parameters(), []);
// Walk the printed lines: sources take their stored payloads, and
// every computed node expresses its opcode over operands computed
// earlier. Under the exact posture, this is the interpreter.
let values = Numerics::exactly(|| {
let mut values: Vec<Tensor<f32>> = Vec::new();
for node in network.nodes() {
let value = if node.is_source() {
network.payload(node.symbol()).expect("sources hold payloads").clone()
} else {
let operands: Vec<&Tensor<f32>> =
node.operands().iter().map(|symbol| &values[symbol.index()]).collect();
node.opcode().express(&operands)
};
values.push(value);
}
values
});
let by_hand = &values[loss.index()];
assert_eq!(by_hand.scalar().to_bits(), oracle.of(loss).scalar().to_bits());
show::lines([
format!("walked {} lines by hand, outside the crate", values.len()),
format!("loss by hand: {by_hand}"),
format!("loss by interpreter: {}", oracle.of(loss)),
"same bits: asserted".to_string(),
])That walk is the seam a new AD mode plugs into. Replace Tensor with a payload that carries a value and a tangent, and the same loop computes a directional derivative; examples/forward_mode.rs does exactly that and grades it, bit for bit, against the reverse scan. AD as a named reading is the story.
Run it#
cargo run --example walkthroughThen change the graph — a second layer, a log_softmax, a bigger batch — and read the six plates again. Or do the same in the playground without leaving this site.