API Entry Points And Runtime Artifacts🔗

Training and a standalone autograd.grad call build different runtime artifacts:

The autograd.* transforms do not consult the trainer's settings. In Autodiff.lean, gradients uses the general vjp path with a unit scalar seed. That path calls lowerToTypedGraph, then vjpChecked, whose compiled implementation uses checked graph execution and saved reverse programs. A proved compiler simplification replaces its raw-tape specification. Scalars use the same lowering as other tensor outputs. The graph describes the reusable computation; each invocation supplies its own values and saved local rules. Higher-order transforms in the same file lower graphs to tapes over dual scalars, so the reverse calculation also carries directional derivatives. The trainer does not use these transforms; it records an eager tape.

Consequently, comparing autograd.grad with an eager training step compares two execution paths. A discrepancy could come from either path or from the state supplied to it. Inspecting only the eager tape would leave the transform's typed graph unexamined.

Runtime Artifact Types🔗

The type parameters reveal which information callers must supply statically and which information the representation stores internally:

Tape : Type → Type

Torch.TypedGraph : (α : Type) → [Storage α] → List Shape → Shape → Type

NN.IR.Graph : Type

NN.Backend.IR.GraphKernelPlan : Type

Tape V takes the type of the values it stores and nothing else: the CPU stores Spec.SomeTensor α, CUDA stores device buffers, and the nodes, parents, and reverse sweep are the same for both. Each stored value still carries its own shape, but the tape's type does not record a context of node shapes. Agreement between a node and the parents named by its identifiers is checked at runtime. That is a deliberate trade, and we come back to it below.

Torch.TypedGraph α Γ τ carries a context Γ : List Shape and an output shape τ. A node cannot refer to a parent at the wrong shape, because the shape is part of the graph's type. This is the same trick the sequential architecture language uses in GraphSpec: One Architecture, Several Meanings, applied one level lower.

NN.IR.Graph : Type has no α. The canonical IR is not parameterized by a scalar type, so the same graph can be interpreted over the reals, over intervals, and over binary32 without being rewritten. See Graph IR.

NN.Backend.IR.GraphKernelPlan : Type has no scalar or static shape indices. It holds node identifiers and capsule metadata, and it contains no tensors whatsoever. A plan cannot witness that anything executed. Accepting a plan is a statement about a choice, not about a computation.

Here is the same set with the two remaining columns filled in:

A CUDA Graph capture is a separate device launch-and-replay mechanism, and TorchLean does not use it. The other boundaries also require explicit connections. An accepted kernel plan needs runtime binding before it says anything about execution, and an IR semantics theorem needs a provider refinement result before it applies to a native tape node.

The four type signatures divide the debugging problem into useful pieces. A wrong tensor value belongs to an execution artifact such as the tape or evaluated graph context. A wrong operation tag belongs to the canonical IR. A surprising provider belongs to the kernel plan and its runtime binding. None of the last two objects contains the floating-point intermediate that caused a loss to become nonfinite. Keeping these objects distinct tells us which additional record to request instead of treating every representation called a graph as interchangeable.

The scalar parameter also has consequences. The same shape-indexed program can be interpreted with different scalar operations, but a tape already contains values of its chosen α. Changing arithmetic means evaluating with the other scalar interpretation, not relabeling those stored values. Shape agreement survives that choice; numerical agreement requires a separate argument.

Manual Tape Construction🔗

Tape V is a grow-only array of nodes. Each operation takes a tape and returns the extended tape with the new node's identifier, inside Except String, so building a tape is a pure computation and its errors are values. To examine the node structure, record two elementwise products sharing the same input, followed by their sum:

z = x \odot x + x \odot x, \qquad \ell = \textstyle\sum_i z_i

Here \odot denotes elementwise multiplication, z has the same shape as x, and \ell reduces it to a scalar. The following code uses x=(3,-1.5):

Except.ok (4, 5)

Five nodes, and the scalar output is node 4. Each operation appends exactly one node and returns its identifier, so identifiers are array indices and reverse-mode traversal is a loop from size - 1 down to 0. Nothing here is hidden state: the tape is a value passed from one call to the next.

Ask the tape what it recorded:

Except.ok #[(some "x", #[]), (some "mul", #[0, 0]), (some "mul", #[0, 0]), (some "add", #[1, 2]), (some "sum", #[3])]

The leaf has no parents. Both multiplications name node 0 twice, once for each factor, which is literally what x * x means. The addition names both products. This printout is the graph, and it came out of the same data structure the trainer uses.

A node also carries the piece we did not print, because it is a closure rather than data:

-- A local rule may emit several contributions to the same
-- parent.
backward : SomeTensor α → Result (Array (Nat × SomeTensor α))

Given an upstream cotangent for this node's value, return one contribution per parent. That signature is the whole local interface of reverse mode. Nothing in it mentions the rest of the graph, which is why local derivative rules can be stated and proved one operation at a time.

The result (4, 5) contains the output node identifier and the tape length. Node identifiers start at zero, so the scalar sum at node 4 is the fifth recorded value. The result is not the loss itself. At the chosen input, each square branch contains [9, 2.25]; adding them gives [18, 4.5], and the scalar sum is 22.5. This calculation supplies a forward reference before inspecting any derivative.

The two parent entries [0, 0] on each mul are intentional. Both operands refer to the same input, but they occupy different argument positions in multiplication. A reverse rule owes one contribution for each position. Deduplicating that parent list without changing the rule would lose half of the derivative of a square. Reusing a value and using it only once are different properties of a graph.

Gradient Accumulation On The Tape🔗

Two product branches run from x to z. Each squaring contributes 2x, and the addition sends the output seed to both of its parents, so the leaf's cotangent is

\bar x = 2x + 2x = 4x.

At x = (3, -1.5) that predicts (12, -6). Seed the scalar with 1 and read node 0:

Except.ok [12.000000, -6.000000]

Each mul returns two contributions to node 0, one for each occurrence of x, giving four contributions in total. If later contributions overwrote earlier ones instead of adding, the result would be wrong. Accumulation happens in exactly one place, SomeTensor.add in Engine/Core/Base.lean, which checks that the two shapes agree before adding. Reverse-mode AD on a DAG is a sum over paths (Baydin et al. (2018)), and this one function is where the sum is taken.

This is also why TorchLean's autograd proofs come in two layers rather than one. There are primitive derivative facts, one per operation, and there is a global soundness statement about the traversal that composes them. A library that proves only the first layer has proved that its rules are right and said nothing about whether they are combined correctly. Proving Autograd Correct is where both layers are stated and discharged.

# Compare two live branches with one branch detached from
# autograd.
import torch

x = torch.tensor([3.0, -1.5], requires_grad=True)
z = (x * x + x * x).sum()
z.backward()
print("both branches live:", x.grad.tolist())

x.grad = None
z = (x * x + (x * x).detach()).sum()
z.backward()
print("one branch detached:", x.grad.tolist())

both branches live: [12.0, -6.0]
one branch detached: [6.0, -3.0]

Stop-Gradient Nodes🔗

detach is often described as "turning off gradients", which suggests a mutable switch. In this engine it is a node like any other, and the node is short enough to read in full. Here is the ordinary Float case of what Session/Recording.lean appends, transcribed into the pure API. The generic session also clears scalar differentiation metadata when the carrier has it.

Except.ok [6.000000, -3.000000]

2x instead of 4x, matching PyTorch's second line to the digit. Three fields did that work: backward returns an empty array, so no contribution flows; parents still names the source, so the forward edge remains visible to anything that reads the graph; and requiresGrad := false causes parent accumulation to skip this node. The forward tensor value is reused without recomputing or copying its payload, so the forward result is unchanged.

The node list makes the difference visible:

Except.ok #[(some "x", #[]), (some "mul", #[0, 0]), (some "mul", #[0, 0]), (some "detach", #[2]), (some "add", #[1, 3]),
  (some "sum", #[4])]

Six nodes now, and the addition's parents are 1 and 3: the live product and the detached copy. The tape still knows the whole forward computation. Only the reverse rule changed.

Except.ok [0, 1, 4, 5]

Differentiation Through The Transform API🔗

We now have a hand-built answer. Compare it with the answer the public transform API gives, which takes an entirely different route: it lowers the program to a typed scalar graph and differentiates that. Write the program once against the operation interface, so that the same text can be interpreted either way:

both branches live = [12.000000, -6.000000]
second detached    = [6.000000, -3.000000]

The tape, the typed graph, the closed-form calculation, and PyTorch agree on 4x and 2x at this input. The comparison checks several implementations of the same branch structure. If only one result changes, that narrows the investigation to its construction or execution path; shared rules can still produce shared errors.

A useful independent reference here is the algebraic calculation above: two square branches give 4x, and the chosen stop-gradient rule on one branch leaves 2x. For a more complicated program, derive a small special case or a directional check before comparing two execution routes that reuse the same rules.

Eager And Graph Interpreters🔗

The reason the same source text can be executed eagerly or recorded into a graph is not a flag inside the runtime. It is two instances of one class. Ops is the operation interface that Torch.mul, Torch.add and the rest are defined against, and the trainer provides two instances of it, in Trainer/EagerOps.lean and Trainer/GraphOps.lean. The difference between them is visible in a single field, Ref, which says what a handle to an intermediate value actually is:

Both equalities hold by reduction, as the rfl proofs show. Under the eager instance a reference is a TensorRef, an identifier into the session's tape, and each operation executes immediately and appends a node. Under the graph instance a reference is a GraphM.Var, an SSA variable, and each operation records a node without computing anything.

This is the tagless-final pattern (Kiselyov (2012)): one object language, several interpreters, with the object language's types carried by the host language. The two fields of a GraphSpec primitive use the same design. A shared operation interface keeps the program text aligned, while the interpreters still need correspondence checks for the operations they implement.

The maintained eager handle also records a session owner and generation. A reset can reuse numeric node 0, so its number alone cannot establish that an old handle belongs to the new recording. The runtime checks those fields before accepting a session reference. This protects a different property from shape safety: two tensors may both have shape [2] while only one belongs to the current execution.

Typed Graph Execution🔗

Typed graph execution records the model's scalar loss once as a typed SSA graph. Nodes carry the forward behavior and the executable rules used for JVPs and VJPs. Each evaluation supplies current parameters and inputs to that recorded graph. A separate typed reference selects the result, so the graph may return an input or an earlier node without adding a dummy final operation, and forward and reverse execution read and seed that exact reference.

The graph's indices enforce shape agreement between node references, but do not establish derivative correctness. The executable node payload holds three functions, forward, jvp and vjp, with runtime validation and optional certified preparation. The preparation certificate preserves the forward, validation, and dense VJP programs; it supplies no derivative-correctness law. The proof-carrying node in Tape/Algebra/Soundness.lean adds the local adjointness law relating its jvp and vjp, and the proof-carrying graph composes those local laws into a graph theorem. Execution can proceed without this adjointness certificate; the graph verification theorem requires it. Proving Autograd Correct follows that ladder up.

The public model-level name specializes the context of this same graph type:

nn.TypedGraphModel is an abbreviation whose context is model state followed by one input tensor. It introduces no additional runtime representation; the equality unfolds to rfl.

Runtime.Autograd.TypedGraph.compileChecked returns a Compiled execution containing an indexed primal context and saved local VJP programs. Compiled.backwardDenseFrom accepts a seed pack for the complete context. The ordinary checked and pure TypedGraph VJP APIs use this path through proved compiler simplifications; the reverse-mode Jacobian transforms in Autodiff.lean use it directly. Compiled.toTape and lowerToTapeChecked retain the raw Tape interface. The accompanying theorems identify the complete checked error/result and every dense gradient under Storage and Add, without assuming associative addition. They preserve the stored executable rules; mathematical derivative correctness and agreement with GPU kernels remain separate obligations.

Saved backward shares uniform contribution histories while preserving each dense scalar addition, including the complete local contribution of a node with repeated parents. Compression uses the certified carrier equality in Storage.decEq?: native Float, Float32, and UInt8 provide it; other storage instances default to dense backward. This equality distinguishes signed zeros but does not require raw NaN sign or payload identity. Custom nodes without certified local preparation or compact contributions retain materializing adapters, and noncompressible histories can require quadratic replay. The raw Tape adapter also retains full prefix contributions.

Construction and final disposal can require quadratic work because the graph representation retains full shape-prefix lists. Efficient reusable execution does not imply a linear graph lifecycle or a general complexity theorem.

Recording and lowering do not imply the optimization, fusion, scheduling, or native code generation associated with a compiler such as torch.compile.

Canonical IR Representation🔗

Typed graph nodes hold Lean functions. Those functions can execute an operation, but they do not expose its structure as inspectable data. NN.IR.Graph instead represents operations such as linear, ReLU, reduction, normalization and shape transforms as data. Each node records an operation tag, parent identifiers, and its output shape. Input shapes are obtained from parent nodes; parameter data is looked up in the separate payload.

Because the operation is a tag, an importer can validate it, a verifier can interpret it, and a code generator can reject what it does not support without executing an arbitrary function. The scalar-independent representation also permits several interpretations: Graph IR builds one small graph and evaluates it over the reals, over intervals, and over IEEE binary32 scalars without touching the graph.

The canonical IR is executed by its own semantics, NN.IR.Semantics, and it is forward-only. It supplies no JVP or VJP rules; differentiating an imported operation would require additional rules and corresponding correctness statements.

Both paths honor the homogeneous-element contract from Tensors And Shapes: a typed graph is parameterized by one α, and an IR denotation receives one payload over α and produces values over that same α.

Shape Erasure And Recovery🔗

A parameter pack holds differently shaped tensors that share one element type:

-- Each state entry keeps its own shape; the pack also fixes
-- the order of the two layers.
[weight1 : Tensor α [8,2],
 bias1   : Tensor α [8],
 weight2 : Tensor α [1,8],
 bias2   : Tensor α [1]]

The pack type preserves that dependent list. A tape cannot, because Tape V has no shape index, and neither can a registry that has to iterate over parameters by name. Both store Spec.SomeTensor α, a structure holding a shape and a tensor indexed by that shape. The collection can therefore hold different shapes while each entry retains its own shape-value relationship.

@Spec.SomeTensor.cast : {α : Type} →
  [inst : Storage α] → {shape : Shape} → (value : Spec.SomeTensor α) → value.shape = shape → Tensor α shape

@Spec.SomeTensor.ofTensor_cast : ∀ {α : Type} [inst : Storage α] (value : Spec.SomeTensor α) {shape : Shape}
  (h : value.shape = shape), Spec.SomeTensor.ofTensor (value.cast h) = value

First, cast takes a proof. Its last explicit argument is value.shape = shape, so there is no way to recover a typed tensor from a packed one without producing evidence that the stored shape is the expected shape. A runtime comparison can supply that evidence, but rfl cannot establish an arbitrary packed tensor's shape:

Application type mismatch: The argument
  rfl
has type
  ?m.11 = ?m.11
but is expected to have type
  value.shape = [2]
in the application
  value.cast ⋯

Second, ofTensor_cast says that casting a packed tensor to the matching shape and repacking it returns the original packed value. Its stored shape and values survive the round trip.

The instantiated runner also holds parameter names, bounded token inputs, RNG state, optimizer memory, mutable model buffers, and the train or eval flag.

The cast signature takes an equality from the stored shape to the requested shape. Its direction matters: it transports the existing tensor; it does not reshape, truncate, pad, or reinterpret the buffer. The roundtrip theorem says that packing a tensor and recovering it with the matching shape returns that tensor. It does not manufacture an equality for arbitrary packed data.

The printed metavariable names are elaborator details; the useful part of the diagnostic is the demanded equality with [2].

Model Mode And Random State🔗

Some operations depend on mode. Dropout samples a mask during training and applies its deterministic inference behavior during evaluation ((Srivastava et al., 2014)); batch normalization may update running statistics during training and use stored statistics at evaluation ((Ioffe and Szegedy, 2015)); other stochastic or stateful layers follow the same pattern. trained.predict uses the retained runtime state in evaluation mode, and a Trainer.Session follows the same rule in a caller-driven loop, with step in training mode and predict and loss in evaluation mode. Mode is not a backend: one eager CUDA runner can switch modes without changing device or provider profile. Choosing How A Model Runs measures the size of the resulting difference, which is the part that matters when you are trying to tell a mode bug from rounding noise.

For the tape, randomness raises a replay requirement. A stochastic forward pass and its reverse pass must use the same random draw, or the reverse pass computes the derivative of a function nobody evaluated. The tape interface makes replay deterministic, but correct capture of the forward mask is still an obligation of the node implementation:

@Tape.backwardScalar : {α : Type} →
  [inst : Storage α] → [Add α] → [One α] → Tape (Spec.SomeTensor α) → ℕ → Result (Std.HashMap ℕ (Spec.SomeTensor α))

No IO: in the pure Lean model, replay depends on the tape and its stored closures rather than an external generator. A closure could still use an incorrect deterministic mask or an incorrect VJP; purity alone does not prove it differentiates the recorded forward pass. Explicit capture and local correctness laws supply that connection when treating reverse mode as a program on immutable data (Pearlmutter and Siskind (2008)). The grow-only array behind it is efficient for the same reason Lean's own runtime can afford immutable structures ((Ullrich and de Moura, 2019)).

This also affects checkpoints. Inference can require persistent buffers, mode, and preprocessing as well as parameters. Reproducing the next training update needs additional state. Replaying that update also needs the generator state, the optimizer state, and the loader position.

The pure backward type still returns Result: lack of IO does not mean that every supplied tape is valid. The traversal can reject an invalid output id, an incompatible seed, or a parent contribution with the wrong shape. Purity describes how the computation obtains its inputs; validation describes which of those inputs it accepts.

The eager reverse engine also distinguishes an unreached node from a reached node with a zero cotangent. It skips unreached local rules and fills missing entries with zeros when a dense result is requested. This matters at singular values: evaluating a disconnected rule can produce an expression such as 0 * (1 / 0), which need not be zero in floating-point arithmetic. The bridge to a traversal that visits every node therefore needs a zero-preservation condition. That is a mathematical obligation about the rules, not just a choice of array or hash-map storage.

Activation Recomputation🔗

Saving model state to a checkpoint and checkpointing an activation graph solve different problems. Activation checkpointing discards intermediate values during the forward pass and recomputes them during the reverse sweep on a schedule that trades arithmetic for memory (Griewank and Walther (2000)). The tape examples here retain their recorded values; they do not demonstrate that memory optimization.

Recomputation must reproduce the forward values used by the local reverse rules. For a stochastic operation, this includes replaying the same random draw. Agreement of the forward function alone does not establish a memory saving or validate a recomputation schedule.

Runtime And Proof🔗

Relating execution to a derivative requires two kinds of correspondence.

The derivative ladder is internal to Lean:

calculus rule for each primitive
  -> semantic local VJP
  -> well-formed graph or tape composition
  -> global reverse result

The refinement ladder leaves Lean:

executable primitive
  -> declared semantic primitive

and it needs one rung per provider actually used by the run.

The first ladder can be climbed entirely inside Lean for a supported abstract graph, and Proving Autograd Correct climbs it: local adjointness laws compose into a graph theorem, the executed dense reverse sweep is identified with that theorem's backpropagation, and over the reals the result is identified with the adjoint of a Fréchet derivative from Mathlib ((Community, 2020)). The second ladder may be a theorem, a sound checked guard, a numerical certificate, or an explicit trusted boundary. Native CUDA, LibTorch, the compiler, the driver, and the hardware need their own refinement evidence or stated trust assumptions.

Relevant proof sources:

• Autograd/Overview.lean;

• Tape/Algebra/Soundness.lean;

• Runtime/Link.

Kernel Selection Records🔗

Before an eager operation executes, the session resolves its backend profile and binds the accepted capsule to a matching runtime handler; a missing provider or handler is an error. Run

# Inspect the capsules selected for a short eager training
# run.
scripts/lake.sh exe torchlean quickstart_mlp \
  --device cpu --steps 1 --seed 2026 --show-backend

and each operation reports one line the first time it is selected. Here is the ReLU entry from the recorded run, rewrapped to fit the page:

  relu: reference.relu provider=reference trust=checked
    vjp=torchlean-tape reduction=n/a

Choosing How A Model Runs reads the full report, tabulates the seven capsules a CPU run selects, and explains the evidence lines and the empty banners. The field vjp=torchlean-tape names which artifact owns the reverse rule for that operation. It connects the tape we built by hand to the selected kernel: the forward value came from the reference.relu provider, and the backward contribution came from TorchLean's own tape node. When a provider owns its own backward instead, the corresponding derivative claim also depends on that provider's reverse implementation.

Two deep dives follow a single operation the whole way down. Backend Selection gives the selection and assurance model, and GPU and CUDA follows an accepted CUDA operation through storage checks, launch, and backward registration.

Tape Inspection In VS Code🔗

The pure API exposes a tape as a value that tests can compare. The editor widgets present its structure visually when the cursor is placed on

-- Inspect recorded nodes and the cotangents associated with
-- them.
#tape_view ...
#tape_grads_view ...
#tape_trace_view ...

The first shows operation nodes and parent edges. The second evaluates a scalar-output reverse pass and annotates gradients. The third exposes traversal order and the final accumulated cotangents. In the branched example above, these views show both products contributing to the leaf. With detach, the second product remains in the forward graph but contributes nothing to the leaf's cotangent. See Interactive Widgets.

Training Through The Trainer🔗

This command makes a reproducible eager training run, with a fixed seed, dataset, and arithmetic:

# Keep seed and data fixed for a two-step run.
scripts/lake.sh exe torchlean quickstart_mlp \
  --device cpu --steps 2 --seed 2026

== Quickstart: simple MLP training (seed=2026, steps=2) ==
target(heldout)    = [0.200000]
untrained(heldout) = [-0.088261]
dataset size = 25
mean_loss(before training) = 0.495227
mean_loss(after training) = 0.392821
steps=2 arithmetic=native scalar=Float32
  loss=0.495227 -> 0.392821
trained(heldout) = [0.019031]

Each step records a fresh tape for the current parameters and the step's sample, then runs it backward. Choosing How A Model Runs compares arithmetic and device choices through their displayed results.

The two-step log reports a dataset mean falling from 0.495227 to 0.392821. The held-out value 0.019031 is a prediction with target 0.2, not another loss measurement.

Runtime Debugging🔗

A gradient discrepancy can originate in the mathematical reference, the recorded program, the state supplied to it, or the provider that executes it. These checks separate those possibilities.

1. Derive a small reference case, such as 4x, including the reduction and output seed.

2. Print the arithmetic semantics, device, and selected capsules, so you know which artifacts are even in play.

3. Inspect the forward tape and check that the branch you care about is reachable. A missing key in the gradient map and a zero value in it mean different things.

4. Check the output cotangent's shape and values. A wrong seed produces a perfectly correct derivative of the wrong scalar.

5. Check detach, train or eval mode, and stochastic state. They affect the derivative path, forward behavior, or random draw used by the same model source.

6. Compare eager and typed graph CPU execution on the same explicit parameters. Disagreement helps locate a boundary; agreement on one case does not rule out a shared bug.

7. Compare a native provider against the IEEE reference path on a small finite case.

8. Compare the gap with an error bound or a more precise reference. A small difference can be a semantic bug, and a large one can come from numerical instability.

Tape Experiments🔗

The hand-built tape also exposes cases that a successful scalar gradient does not distinguish:

1. A third x*x branch adds another 2x contribution. Its mul node names node 0 twice, just as the first two do.

2. Setting the leaf's requiresGrad := false prevents retention of its cotangent, so rtLeafGrad reports a missing gradient.

3. Moving rtDetachNode from b to a preserves the reported cotangent because the products are equal, but changes the node identifiers on the surviving path.

4. A scalar output independent of the leaf leaves node 0 absent from the gradient map.

5. A mul node with incompatible operand shapes makes Tape.mul return an error value in Except String.

6. Changing a constant only in rtQuad makes the typed graph represent a different function. Comparing its result with the unchanged tape then detects a construction mismatch.

Before defining the canonical IR we cross one more concrete boundary: a named parameter payload moving between PyTorch and Lean, in Exchanging Models With PyTorch. The graph chapters then give that exchanged computation an inspectable semantics, and relate runtime approximation and verification claims to it.