The autograd tape

One tape serves every device. A Tape V is a grow-only array of nodes whose forward values have type V: the CPU stores Spec.SomeTensor α, CUDA stores LibTorch.SomeBuffer. Each node keeps its parents and a local VJP, and the reverse sweeps below walk node ids from last to first, summing the contributions that several children send to one parent.

The dense sweeps are shared between devices. GradValue supplies device-specific gradient checks, addition, and buffer ownership; the CPU hooks are pure tensor operations. CUDA training also uses the sparse LibTorch.Tape.backwardSparse traversal in NN.Runtime.Autograd.Engine.LibTorch.Tape, which releases unretained gradients in IO.

Two traversals live here, and it matters which one a caller runs:

• backwardDense (and its totalized form backwardDenseAll) keeps an Option per node and runs a node's VJP only when that node has received a cotangent, so disconnected nodes are never visited. Skipping is deliberate: on Float, feeding a synthetic zero cotangent through the VJP of a singular value can produce NaN via 0 * (1/0).
• backwardDenseFrom starts from a total gradient array and runs every node's VJP. It is the variant the proof layer reasons about directly (Proofs.Autograd.Algebra.Graph.backwardDenseFrom_lowerGraphToTape_eq_backpropAllCtx).

The two agree whenever every VJP on the tape sends a zero cotangent to zero contributions of the parents' shapes; that is Proofs.Autograd.Algebra.Graph.ZeroPreserving in NN.Proofs.Autograd.Runtime.Link.BackwardDense.

@[reducible, inline]

abbrev Runtime.Autograd.Result (α : Type) : Type

Runtime error monad for the autograd tape.

We use plain Except String (instead of IO exceptions) so the tape constructors remain pure and easy to test. Front-ends that prefer exceptions can use Lean's IO.ofExcept.

class Runtime.Autograd.GradValue (V : Type) : Type 1

What the reverse sweep needs from the values a tape stores.

A gradient for a node has the node's shape. retag re-expresses a gradient whose shape has been checked against a value, check performs any further device check (CUDA compares dtypes and buffer lengths), and add sums two gradients of one node. Backward rules return fresh contributions, so a node keeps its first contribution as it is; drop retires one that no node keeps, which on CUDA releases the device buffer at once.

• shape : V → Spec.Shape

  Runtime shape of a value.

• retag (g v : V) : shape g = shape v → V

  g, given the shape of v, once the two shapes are known to agree.

• check (g v : V) : Result Unit

  Device checks of a gradient g against the value v it belongs to.

• add : V → V → Result V

  Sum two gradients of one node.

• drop {β : Type} : V → β → β

  Retire a contribution that no node keeps, then continue with rest.

class Runtime.Autograd.GradZero (V : Type) : Type

A zero gradient of the same shape as a value, for nodes the sweep never reached.

• zerosLike : V → V

  A zero gradient shaped like v.

structure Runtime.Autograd.Node (V : Type) : Type

One recorded computation step.

value is the forward value, parents the ids of its inputs, and backward its local VJP: given the upstream cotangent of value, it returns (parentId, contribution) pairs. When several children contribute to one parent, the sweep sums the contributions with GradValue.add.

PyTorch comparison: an autograd Function instance with its saved tensors, storing the backward closure directly.

• name : Option String

  Optional node name used for debugging and pretty-printing.

• value : V

  Forward value computed at this node.

• ownsValue : Bool

  Whether this node owns its forward value. On CUDA, shape views and detached references borrow a parent's buffer, and releasing a borrowed value would also invalidate the parent.

• requiresGrad : Bool

  Whether reverse-mode propagation should visit this node.

  If false, reverse-mode traversal skips this node and does not accumulate gradients into it. Operations inherit the disjunction of their differentiable parents' flags; leaves retain the caller's explicit flag.

• parents : Array ℕ

  Parent node ids (dependencies) in the tape.

• cleanup : Array V

  Forward workspace kept only because backward may need it. CUDA releases these buffers when the tape is reset, so long training loops do not wait on finalizers for large intermediates.

• backward : V → Result (Array (ℕ × V))

  Local VJP rule for this node.

structure Runtime.Autograd.Tape (V : Type) : Type

Autograd tape: a grow-only array of nodes.

Node ids are array indices (Nat). All operations append exactly one node and return its id, so reverse mode traverses ids in reverse order.

• nodes : Array (Node V)

  Tape nodes in evaluation order (id = index).

def Runtime.Autograd.Tape.empty {V : Type} : Tape V

Empty tape (no nodes).

def Runtime.Autograd.Tape.size {V : Type} (t : Tape V) : ℕ

Number of nodes stored in the tape.

def Runtime.Autograd.Tape.getNode? {V : Type} (t : Tape V) (id : ℕ) : Option (Node V)

Read a node by id (returns none if out of bounds).

def Runtime.Autograd.Tape.getValue? {V : Type} (t : Tape V) (id : ℕ) : Option V

Read just the stored forward value for a node id.

@[inline]

def Runtime.Autograd.Tape.addNode {V : Type} (t : Tape V) (node : Node V) : Tape V × ℕ

Append a node and return its id.

Invariant: the returned id is t.size, the pre-append size of the tape.

@[simp]

theorem Runtime.Autograd.Tape.addNode_id {V : Type} (t : Tape V) (node : Node V) : (t.addNode node).2 = t.size

addNode returns the current tape size as the fresh node id.

@[simp]

theorem Runtime.Autograd.Tape.size_addNode {V : Type} (t : Tape V) (node : Node V) : (t.addNode node).1.size = t.size + 1

Appending a node increases the tape size by one.

Reverse sweeps

def Runtime.Autograd.Tape.addGradDense {V : Type} [GradValue V] (t : Tape V) (grads : Array (Option V)) (id : ℕ) (g : V) : Result (Array (Option V))

Add one parent contribution into the optional gradient array.

The entry is none until the sweep first reaches the node. A first contribution is kept as it is; later ones are summed in, in the order the sweep produces them.

def Runtime.Autograd.Tape.backwardDenseStep {V : Type} [GradValue V] (t : Tape V) (acc : Array (Option V)) (id : ℕ) : Result (Array (Option V))

Visit one node in the optional-gradient reverse sweep.

Unreached nodes are skipped without running their VJP; reached contributions accumulate in order.

def Runtime.Autograd.Tape.backwardDense {V : Type} [GradValue V] (t : Tape V) (outId : ℕ) (seed : V) : Result (Array (Option V))

Reverse-mode backpropagation producing a dense array of optional gradients.

• The result array has length t.nodes.size.
• Entry id is some g if the node was reached from outId during reverse traversal, otherwise none.
• When multiple paths contribute to the same node, the contributions are summed.
• A node's VJP runs only if the node was reached; see the module docstring for why.

It is loosely analogous to PyTorch's autograd engine walking the dynamic graph and accumulating .grad, but it keeps gradients for every node id rather than for leaves alone. That makes the runtime easier to debug and gives proof-bridge code direct access to intermediate cotangents.

def Runtime.Autograd.Tape.addGradAll {V : Type} [GradValue V] (t : Tape V) (grads : Array V) (id : ℕ) (g : V) : Result (Array V)

Add one parent contribution into a total gradient array.

This is used by the proof-friendly backwardDenseFrom* variants, which start from an explicit gradient for every node.

def Runtime.Autograd.Tape.backwardDenseFromStep {V : Type} [GradValue V] (t : Tape V) (acc : Array V) (id : ℕ) : Result (Array V)

One reverse-mode step at a single node id, updating a total dense gradient array.

Precondition by convention: acc has one entry per tape node, and every entry has the matching node shape. The function checks those conditions dynamically and returns an error if a caller violates them, which makes it the small proof-friendly step used by backwardDenseFromLoop.

def Runtime.Autograd.Tape.backwardDenseFromLoop {V : Type} [GradValue V] (t : Tape V) : ℕ → Array V → Result (Array V)

Reverse-mode accumulation over the first n nodes in reverse order.

The recursion visits node ids n-1, n-2, ..., 0. Passing n = t.nodes.size therefore traverses the entire tape.

def Runtime.Autograd.Tape.backwardDenseFrom {V : Type} [GradValue V] (t : Tape V) (grads0 : Array V) : Result (Array V)

Reverse-mode accumulation starting from an explicit dense gradient array.

This always runs every node's VJP (in reverse order) and keeps a gradient for every node id. The eager runtime does not call it; it calls backwardDenseAll, which agrees with this function on zero-preserving tapes (Proofs.Autograd.Algebra.Graph.backwardDenseAll_eq_backwardDenseFrom).

def Runtime.Autograd.Tape.totalize {V : Type} [GradZero V] (t : Tape V) (grads : Array (Option V)) : Array V

Fill in the gradient of every node the sweep did not reach with a zero of the node's shape. Reached nodes keep their computed gradient.

def Runtime.Autograd.Tape.backwardDenseAll {V : Type} [GradValue V] [GradZero V] (t : Tape V) (outId : ℕ) (seed : V) : Result (Array V)

Reverse-mode accumulation that returns a gradient for every node id.

Propagation uses backwardDense, so local VJP closures run only for nodes reached from outId. Unreached nodes then receive explicit zeros. This is necessary at singular forward values: applying a disconnected VJP to a synthetic zero cotangent can manufacture NaN through expressions such as 0 * (1 / 0), even though the mathematical gradient of the selected output with respect to that node is zero.

On zero-preserving tapes (in particular every tape produced by lowerGraphToTape) it returns exactly what backwardDenseFrom returns from the one-hot seed array; see NN.Proofs.Autograd.Runtime.Link.BackwardDense.