Crate topos

Expand description

topos is an autodiff compiler stack. Record a graph, inspect it, differentiate it, compile it, emit it. The spec is an immutable Network; the state is a caller-owned Parameters.

Expressions record onto a Tape. Sealing yields the network. forward materializes every value, backward differentiates one scalar target, and step is a pure data transform of the parameters. Training never touches the graph:

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

// Record the graph in one closure; the return value is the set
// of names that leave the tape, detached to symbols in one call.
// Operators record as they run; values are `Copy` and never
// consumed. A scalar is a rank-0 tensor: the graph is always
// tensors, and the element type (`f64` here) is the open seam.
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();

// The graph is recorded once; every step feeds one sample of the line
// `y = 2 * x` and steps the parameters, leaving the network untouched.
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)
    });
}

let learned = parameters.of(w).scalar();
assert!((learned - 2.0).abs() < 1e-6);

Differentiation comes in a hierarchy of three, in this order of recommendation. Tape::differentiate records the chain rule as ordinary nodes and answers Adjoints — the derivative as spec: lower a forward-only entry over adjoints.roots() and fusion and liveness apply to the chain rule itself, with Run::recorded_gradients bridging to step. Run::backward (the loop above) is the interpreter applying the same rules without recording — the oracle the transform is proven against bitwise, shipped forever. Entry::backward is neither: a memory posture that retains what the engine scan reads, so a plan that did not record its derivative can still answer backward.

§The stack: one spec, named interpretations

The tape is the spec; everything after it is a derived interpretation of the same columns, each with a printable artifact, and the whole compiler is this list:

spec       Tape / Network        record, describe
shape      inferred at record    panics at the recording expression
value      BoundEntry::interpret the oracle; Network::forward is the whole-spec form
cotangent  Run::backward         the engine reverse scan, oracle of reverse mode
trace      Tape::differentiate   the same rules recording themselves (Trace)
schedule   BoundEntry::lower     Plan: keep-set, liveness, election; describe
catalog    Plan::patterns        elected offers as data, never rewrites
text       Plan::emit_stablehlo  the interchange boundary

The value and cotangent rows compute over Tensor; the trace row records over Trace — one derivative-rule body, two interpretations of the recordable vocabulary (Recordable). A new idea plugs in at a named seam, costed like an opcode: an element type at Element, a transcendental at MapOperation, a fusion as a pattern plus matcher, an AD mode as a recording interpretation proven against Run::backward, an industrial target as an emission sibling consuming Plan. The core stays closed; the table is how the crate refuses a pass manager and still says yes to research.

§Two surfaces, one crate

Two audiences read this crate, and each has a map — rustdoc modules that only re-export, so use topos::Tape keeps working and nothing moves:

model — write a network, train it, checkpoint it: the recording and run types, the neural facades, the optimizers.
compiler — inspect, lower, emit, extend: the printable IR (Opcode, Node, describe), the catalog as data, the recording interpretation (Trace), the element seam, the backend interrogation types, and the reference kernels.

// The compiler surface in three lines: print the spec, lower an
// entry, emit the schedule.
println!("{}", network.describe());
let plan = network.entry([loss]).lower();
println!("{}", plan.emit_stablehlo().expect("every operation lowers"));

Modules§

checkpoint: Module checkpoints: capturing a module tree’s parameter payloads and restoring them into a Parameters state.
compiler: The compiler surface: inspect, lower, emit, and extend the stack.
init: Deterministic initializer factories for neural building blocks.
model: The model surface: write a network, train it, checkpoint it.
reference: The bitwise reference kernels, published for differential testing.

Structs§

Adam: The Adam optimizer: gradient descent through bias-corrected first and second moment estimates (Kingma and Ba, 2015).
AdamW: Adam with decoupled weight decay (Loshchilov and Hutter, 2019): the same moment machinery, with learning_rate * decay * parameter subtracted directly from the parameters the policy selects.
Adjoints: The recorded reverse-mode result: one gradient symbol per wrt entry, in wrt order, paired with the entry it differentiates.
BatchNorm: A batch-normalization layer over [batch, features] values (Ioffe & Szegedy, 2015): every feature is standardized and passed through the learned per-feature affine scale * normalized + shift.
BatchNormTask: One training-mode batch normalization over a contiguous [batch, features] buffer: the fused form of the recorded formula — center by the batch mean, scale by the epsilon-stabilized deviation, apply the learned affine — offered to the backend chain as a single task.
Bf16: A brain-float 16 payload: the top half of an f32, one sign bit, eight exponent bits, and seven stored mantissa bits.
BoundEntry: An Entry bound to the network that will execute it: the builder the common road goes through, and the type that carries the two executor verbs.
Conv2d: A 2-D convolution layer over [batch, channels, height, width] values: the conv2d formula with its kernel stack and bias held as parameters.
Dropout: A mask-fed dropout: the expression multiplies its input by a declared mask input whose default payload is all ones, so an unfed run is the identity — inference is the absence of a feed, not a mode.
Entry: A function exported from a network: the declared reading — roots, observes, memory posture, numerics — that every executor takes.
Field: A value-aligned buffer over the nodes of one network’s recording.
GemmTask: One dense matrix-multiplication job: m x k times k x n, each operand a spanning slice read through two strides.
LayerNorm: A layer-normalization layer over [batch, features] values (Ba, Kiros & Hinton, 2016): every sample is standardized by its own feature statistics and passed through the learned per-feature affine scale * normalized + shift.
Linear: The affine transform input.matmul(weights) + bias, unfused: an activation is its own composition stage, which unlocks the orderings a bundled activation forbids (pre-norm blocks, activation-before-projection).
MapTask: One whole-buffer elementwise transcendental as an offerable task: a MapOperation paired with its elements, the map chains’ twin of GemmTask.
Mlp: A multilayer perceptron: affine stages chained by topology, the convenience constructor over Linear and Activation.
Network: The sealed phase of a recording: an immutable computation-graph spec.
Node: One recorded node of the public IR view: a Copy-cheap snapshot of the columns — name, operands as Symbols, inferred shape — detached from the tape, so a node outlives locks and phases.
Normalization: A recorded batch-normalization expression: the output together with the batch statistics it normalized by.
Normalized: A batch-normalization task’s whole product: the normalized output with the batch statistics it normalized by, mirroring the recorded formula’s root and named results.
Parameters: A payload per parameter slot of one network family: the live weights, and every other table aligned to them.
Path: The full structured path of one parameter in a module tree, ending in the parameter’s own name. Display renders the conventional dotted form (blocks.0.attention.query.weights) for humans and format adapters.
PatternMatch: One recognized pattern as data: its kind, the root node whose result is the group’s result, and every claimed node.
Plan: A compiled lowering of a recorded graph prefix: which nodes a run must evaluate, which values the caller may read, and which buffers may be freed the moment their last consumer has run.
RmsNorm: A root-mean-square normalization layer over [batch, features] values (Zhang & Sennrich, 2019): every sample is divided by the root mean square of its own features and scaled per feature.
Run: The materialized payloads of one forward run.
Sequential: An ordered chain of modules: each stage’s output feeds the next.
Sgd: Plain stochastic gradient descent: the strategy every example’s hand-written loop applies, as the trait’s simplest implementation — stateless, so the struct is a unit, and every richer optimizer is this plus state.
Shape: The runtime extent of a payload along each axis, outermost first.
Symbol: A detached, Copy name for a recorded value: the currency of every phase after recording.
Tape: The construction phase of a network: an append-only record of every node of one computation graph.
Tensor: A dense tensor with an immutable, runtime-defined Shape and a shared element buffer read through a strided layout.
Trace: A payload that records instead of computing: the second interpretation of the derivative rules.
Value: A Copy proxy to a value recorded on a Tape: the operand of recording.

Enums§

Activation: The nonlinearity applied to a neural building block’s affine output.
Backend: The implementers of named formulas: everything that can serve work faster than the reference paths, in LLVM’s sense of the word — hardware kernel providers, the crate’s own fused kernels, and the translation library alike.
BackendUnavailable: Why a Backend would decline all work in this build.
Coverage: One backend’s coverage of one formula: whether it has a kernel, and under what terms.
Dispatch: How a backend’s kernels are reached: the execution-context attribute that replaced the home/abroad dichotomy.
EmitError: Why a plan declined to emit.
Fidelity: The certified fidelity a kernel meets against the oracle.
Formula: Every formula the acceleration stack knows by name: the one vocabulary, leaf and composed entries together.
MapOperation: One unary elementwise transcendental: the shared vocabulary of the IR’s Map node and the backend chain’s whole-buffer map task, mirroring GemmTask for the operations whose scalar form is a libm call the compiler cannot vectorize.
Numerics: The numerics posture of an execution scope: the fidelity it demands of every kernel.
Opcode: The public opcode of one recorded node: the payload-free twin of the engine’s operation enum.
PatternKind: The closed set of graph patterns the catalog recognizes: the public kind of a PatternMatch.
Precision: The seam’s forwarding precisions: the element types that route payload tasks to hardware kernels.
Segment: One step of a parameter path: a child’s position or a field’s static name.

Traits§

Detach: The one-call mass detach: recording-phase handles become the names later phases speak.
Differentiable: The base element contract: the arithmetic of a number that can fill a Tensor.
Element: A number that can fill a Tensor: the open payload seam.
Elementary: Elementary numeric functions of an element, plus its backend hooks.
Emittable: An element type StableHLO emission can render: its MLIR type name and the literal forms MLIR’s float syntax accepts.
Module: A named, parameterized recording function: the unit of model composition.
Optimizer: A training-step strategy: how gradients and the current parameter state become the next state.
Recordable: The recordable vocabulary: the operations derivative rules — and any payload-generic algorithm — are written against.
Visitor: The traversal callback of Module::visit: a path-segment stack plus a parameter sink. Concrete walkers — parameter collection, checkpoint save and restore — are implementations.

Functions§

concat: Records the concatenation of values along axis and returns a proxy to it: each value is padded with zeros to the combined extent at its running offset, and the pads are summed.
conv2d: Records the 2-D convolution of input by weights plus bias on their network and returns the [batch, filters, out_height, out_width] output value.
cross_entropy: Records the cross-entropy loss of logits against targets on their network and returns the rank-0 loss value.
max_pool: Records the size x size max pooling of the [batch, channels, height, width] value input with stride and returns the pooled [batch, channels, out_height, out_width] value.
named_parameters: Returns every parameter in module’s tree with its structured path, in visit order: the name map of the serialization boundary.
parameters: Returns the symbols of every parameter in module’s tree, in visit order.
stack: Records the stacking of values along a new axis at axis and returns a proxy to it: each value gains an extent-1 axis there (unsqueeze) and the lifted values concatenate.

Type Aliases§

Gradients: The gradients of one backward run: the derivative of the run’s target with respect to every node.

Crate topos

Crate topos Copy item path

§The stack: one spec, named interpretations

§Two surfaces, one crate

Modules§

Structs§

Enums§

Traits§

Functions§

Type Aliases§

Crate topos