Engine Core

A small dynamic (DAG) autograd engine.

This is intended to be the "runtime" counterpart to Spec.OpSpec: instead of manually writing backward passes end-to-end, you build a tape during a forward pass and then call backward.

Design goals:

• Works for arbitrary Tensor α s shapes using Spec.SomeTensor values.
• Supports DAGs (shared subexpressions): gradients are accumulated by summation.
• Preserves typed tensors at operator boundaries and reports dynamic shape mismatches as errors.

This module defines and executes the tape; it does not by itself prove every registered VJP. Analytic correctness theorems for supported operations live under NN.Proofs.Autograd, while runtime checks cover the integration between those operations and the tape. For the imperative frontend over this engine, see NN.Runtime.Autograd.Torch.Core.

References (PyTorch / background reading):

• PyTorch docs: torch.autograd and "Autograd mechanics": https://pytorch.org/docs/stable/autograd.html https://pytorch.org/docs/stable/notes/autograd.html
• "micrograd" (small autograd engine, useful for intuition): https://github.com/karpathy/micrograd

Core Types

The eager autograd engine is built out of a few small pieces:

• Result: a pure error monad used throughout the tape API.
• Spec.SomeTensor: a shape-erased tensor used to store heterogeneous node values on one tape.
• Node: one recorded computation step, with a local backward (VJP) rule.
• Tape: a grow-only array of nodes; reverse-mode traversals walk it in reverse order.

@[reducible, inline]

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

Runtime error monad for the eager autograd engine.

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 okOrThrow.

def Runtime.Autograd.okOrThrow {α : Type} : Result α → IO α

Convert an Autograd.Result into an IO action by throwing IO.userError on failure.

This is mainly used by the imperative Torch/TorchLean front-ends to keep their code readable.

def Runtime.Autograd.SomeTensor.add {α : Type} [TorchLean.Storage α] [Add α] (a b : Spec.SomeTensor α) : Result (Spec.SomeTensor α)

Accumulate two packed tensors by elementwise addition, with a dynamic shape check.

This is the heart of DAG support: if two different paths contribute gradients to the same parent, we sum the contributions.

structure Runtime.Autograd.Node (α : Type) [TorchLean.Storage α] : Type

A tape node representing a single tensor value in the recorded computation graph.

Fields:

• value: the forward value (shape-erased).
• parents: ids of parent nodes in the tape.
• backward: a local VJP rule. Given an upstream cotangent for value, it returns a list of (parentId, parentCotangent) contributions (one per parent, usually).

PyTorch comparison: analogous to an autograd Function instance + saved tensors, but here we store the backward closure directly.

• name : Option String

  Optional node name used for debugging and pretty-printing.

• value : Spec.SomeTensor α

  Forward value computed at this node (shape-erased).

• 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.

• parents : Array ℕ

  Parent node ids (dependencies) in the tape.

• backward : Spec.SomeTensor α → Result (Array (ℕ × Spec.SomeTensor α))

  Local VJP rule for this node.

  Given an upstream cotangent for value, return an array of (parentId, parentCotangent) contributions. If multiple children contribute to the same parent, the engine will sum contributions via SomeTensor.add.

structure Runtime.Autograd.Tape (α : Type) [TorchLean.Storage α] : Type

Autograd tape: a grow-only array of nodes.

Node ids are array indices (Nat). All ops append exactly one node and return its id. This makes it easy to implement reverse-mode by traversing ids in reverse order.

• nodes : Array (Node α)

  Tape nodes in evaluation order.

  Node ids are array indices (Nat). Each tape op appends exactly one node and returns its id.

Tape Construction

The Tape namespace provides pure constructors for building a recorded computation graph. Each op appends exactly one node and returns the updated tape plus the new node id.

If you prefer an implicit tape-threading style, see NN.Runtime.Autograd.Engine.TapeM.

def Runtime.Autograd.Tape.empty {α : Type} [TorchLean.Storage α] : Tape α

Empty tape (no nodes).

def Runtime.Autograd.Tape.size {α : Type} [TorchLean.Storage α] (t : Tape α) : ℕ

Number of nodes stored in the tape.

def Runtime.Autograd.Tape.getNode? {α : Type} [TorchLean.Storage α] (t : Tape α) (id : ℕ) : Option (Node α)

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

def Runtime.Autograd.Tape.getValue? {α : Type} [TorchLean.Storage α] (t : Tape α) (id : ℕ) : Option (Spec.SomeTensor α)

Read just the stored forward value for a node id.

def Runtime.Autograd.Tape.addNode {α : Type} [TorchLean.Storage α] (t : Tape α) (node : Node α) : Tape α × ℕ

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 {α : Type} [TorchLean.Storage α] (t : Tape α) (node : Node α) : (t.addNode node).2 = t.size

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

@[simp]

theorem Runtime.Autograd.Tape.size_addNode {α : Type} [TorchLean.Storage α] (t : Tape α) (node : Node α) : (t.addNode node).1.size = t.size + 1

Appending a node increases the tape size by one.

def Runtime.Autograd.Tape.leaf {α : Type} [TorchLean.Storage α] {s : Spec.Shape} (t : Tape α) (value : TorchLean.Tensor α s) (name : Option String := none) (requiresGrad : Bool := true) : Tape α × ℕ

Add a leaf node (no parents).

PyTorch comparison: a tensor that enters the graph as a leaf (e.g. input or parameter value).

def Runtime.Autograd.Tape.requireValue {α : Type} [TorchLean.Storage α] {s : Spec.Shape} (t : Tape α) (id : ℕ) : Result (TorchLean.Tensor α s)

Read a typed tensor value from a tape node id.

This is the main "dynamic check" boundary in the eager runtime:

• fails if the id is invalid, or
• fails if the stored runtime shape doesn't match the expected dependent shape s.

def Runtime.Autograd.Tape.requireGrad {α : Type} [TorchLean.Storage α] {τ : Spec.Shape} (dLdyAny : Spec.SomeTensor α) : Result (TorchLean.Tensor α τ)

Read a typed upstream gradient tensor from a packed tensor.

This is the backward analogue of Tape.requireValue: it checks that the upstream gradient has the expected shape τ and then performs the dependent cast.

def Runtime.Autograd.Tape.unary {α : Type} [TorchLean.Storage α] {σ τ : Spec.Shape} (t : Tape α) (opName : String) (xId : ℕ) (forward : TorchLean.Tensor α σ → TorchLean.Tensor α τ) (backward : TorchLean.Tensor α σ → TorchLean.Tensor α τ → TorchLean.Tensor α σ) : Result (Tape α × ℕ)

Generic constructor for unary ops.

You provide:

• forward : Tensor α σ → Tensor α τ
• backward : Tensor α σ → Tensor α τ → Tensor α σ (a VJP rule; note it may depend on the input)

The returned node stores the forward value and a backward closure that checks the upstream gradient's shape and returns the parent contribution.