3.3.1. PyTorch Mental Model🔗

PyTorch's default workflow is roughly:

1. create parameters and modules,

2. run a forward pass,

3. let autograd record a dynamic computation graph for the ops run during the forward pass,

4. call loss.backward() to accumulate gradients into tensor .grad fields,

5. call optimizer.step() to mutate the parameters,

6. optionally zero the gradients and repeat.

That is the mental model we use here, because TorchLean deliberately mirrors the shape of that workflow while changing the semantics of the bookkeeping.

TorchLean mapping:

• PyTorch parameters and modules become explicit typed values and parameter bundles.

• PyTorch's dynamic autograd tape corresponds to NN.Runtime.Autograd.Engine.Core.Tape.

• PyTorch .grad accumulation becomes explicit gradient results returned by reverse mode.

• PyTorch optimizer.step() corresponds to train.step, train.stepper, or session level update helpers.

• PyTorch eager execution corresponds to TorchLean eager mode.

• PyTorch torch.compile and graph lowering intuition correspond to TorchLean compiled mode and the shared IR.

The key difference is that TorchLean keeps the graph and the shapes visible in the types instead of burying them in mutable runtime objects.

Two more layers are worth keeping in mind for the rest of the runtime chapter:

• NN.Runtime.Autograd.TorchLean.Module helps cast Float literals into the active runtime scalar and package a scalar-loss model in a way that is easier to instantiate repeatedly.

• NN.Runtime.Autograd.TorchLean.Session is the explicit training environment: it owns leaves, tapes, RNG state, and backend selection. The workflow feels familiar from PyTorch, but parameters and gradients stay explicit.

Those are the pieces that make TorchLean's runtime feel familiar without hiding the bookkeeping.

3.3.2. Runtime Design Decisions🔗

The runtime layer is deliberately not a clone of PyTorch internals. It borrows the workflow people know, then changes the representation so the computation remains inspectable from Lean.

1. Gradients are data, not hidden tensor fields. The reverse pass returns gradient values rather than mutating .grad slots behind the scenes.

2. Tapes are explicit. Eager execution records values, parent ids, and VJP closures in a Lean data structure that widgets and proof bridges can inspect.

3. Compiled execution is separate from eager debugging. Compiled mode gives a stable SSA/DAG-like artifact for repeated evaluation and proof alignment, while eager mode stays close to step by step debugging.

4. The verifier uses an IR with operation tags, not arbitrary closures. Runtime closures are great for execution, but verification needs explicit operation tags; that is why NN.IR.Graph exists.

5. CUDA is a backend choice, not a change in model meaning. Device buffers accelerate supported float32 operations, while the Lean spec and IR semantics remain the reference meaning.

This is the main tradeoff of TorchLean's runtime: it is more explicit than a Python training loop, but that explicitness is what lets examples, widgets, proof statements, and verification tools talk about the same computation.

import torch

x = torch.tensor([0.5, -1.2], requires_grad=True)
y = (x * x).mean()
y.backward()
print(x.grad)

import NN

open Spec
open Tensor
open NN.Tensor
open NN.API

def sumsq : autograd.fn1.Fn (Shape.Vec 2) Shape.scalar :=
  fun x => do
    let y ← nn.functional.square x
    nn.functional.mean y

def demo : IO Unit := do
  let x : Tensor Float (Shape.Vec 2) := tensorND! [2] [0.5, -1.2]
  let g ← autograd.fn1.grad (α := Float) sumsq x
  IO.println s!"grad = {Spec.pretty g}"

3.3.3. What the Runtime Adds🔗

The runtime is where TorchLean stops being only a collection of definitions and theorems and starts behaving like something that can be run and inspected. That matters for two rather different audiences:

• debugging a model, loss, or training loop,

• and connecting executable artifacts to a soundness theorem.

The key design principle is that the runtime is not a second source of truth. It is a controlled execution layer whose artifacts should continue to align with the spec and IR semantics described elsewhere in the manual.

This is why TorchLean splits the system into layers:

• the spec layer says what the model means;

• the runtime layer says how that meaning runs and how gradients are produced;

• the compiled/IR layer gives a replayable artifact that verification can inspect;

• the proof layer states theorems about the same underlying object, not a separately exported copy.

That is different from the usual Python stack, where the executed object, the exported object, and the verified object often live in different toolchains.

3.3.4. Two Execution Modes🔗

TorchLean exposes one front-end with two execution backends:

• Eager: tape recording and reverse-mode backprop in the style familiar from PyTorch.

• Compiled: record a graph-like artifact and run it as a replayable, optimizer-friendly execution.

Many of the curated demos accept --backend eager|compiled.

\text{NN.API program} \;\to\; \{\text{eager tape},\text{compiled GraphData}\} \;\to\; NN.IR.Graph \;\to\; \{\text{IBP},\text{CROWN},\text{certificates}\}

3.3.5. What "eager vs compiled" means (in concrete terms)🔗

Both backends execute the same public model and loss definitions from NN.API source tree. The difference lies in the runtime representation of the computation that was executed:

• Eager builds a Tape of runtime nodes as it executes.

  • Each node stores a value plus a backward closure that knows how to send gradients to parents.

  • This is the closest analogue TorchLean has to everyday PyTorch eager mode.

• Compiled builds a typed SSA/DAG representation (GraphData) and evaluates it in one shot.

  • Nodes are typed records containing forward, jvp, and vjp.

  • This is the "compile once, run many times" backend.

  • It is also the backend that is easiest to connect to autograd correctness proofs, because the graph representation is typed and structurally recursive.

Typical usage split:

• eager for stepping through model code, gradients, or training loops.

• compiled for repeated evaluation speed and a stable graph artifact, within the supported operator set of the compiled backend.

The operational difference shows up by running the same demo twice:

lake env lean --run NN/Examples/Quickstart/SimpleMlpTrain.lean -- --steps 50 --dtype float --backend eager
lake env lean --run NN/Examples/Quickstart/SimpleMlpTrain.lean -- --steps 50 --dtype float --backend compiled

With matching seeds and supported operators, the forward mathematics agree; what changes is mostly how much structure can be inspected step by step:

• eager is easier to step through,

• compiled is easier to replay and connect to proof artifacts.

3.3.6. Runtime Contexts And Named Values🔗

Spec tensors are dependently typed by shape (Tensor α s), but realistic training loops need registries:

• parameter maps ("w1" ↦ weights),

• gradient maps ("w1" ↦ dL/dw1),

• "named values" for debugging.

TorchLean therefore uses an existential wrapper:

Runtime.AnyTensor α

which pairs a Shape with a corresponding tensor, allowing heterogeneous registries.

See:

• NN.Runtime.Context API

• NN.Widgets.Runtime.Context for the infoview viewer.

This design preserves the strong typing of the spec layer while still supporting runtime-style tooling.

3.3.7. A Small Runtime Walkthrough🔗

For a single training step, the objects move like this:

1. Build a model and parameter bundle through NN.API.

2. Run a forward pass with a chosen backend (eager, compiled, or CUDA backed eager where supported).

3. Produce a scalar loss.

4. Run reverse mode to obtain explicit gradients.

5. Pass parameters and gradients to an optimizer update.

6. Log metrics and, when needed, inspect the tape or compiled graph.

None of those steps changes what the model means. They only change which runtime artifact is produced and how much structure is available for inspection.

3.3.8. Autograd Tape (Eager Mode)🔗

The eager engine records operations into a Tape:

• nodes store forward values and a local VJP rule (backward closure),

• reverse-mode traverses ids backwards and accumulates gradients.

In PyTorch terms, this is the part of the system that most closely resembles the autograd engine behind loss.backward(), except TorchLean makes the tape explicit and the gradient flow explicit. There is no hidden .grad mutation on tensor objects; the reverse pass produces the gradients directly.

The tape structure is useful because it mirrors the conceptual graph you already see in PyTorch, but it keeps the graph available for debugging and verification instead of making it an implementation detail.

Core implementation:

• NN.Runtime.Autograd.Engine.Core API

import NN

open Spec
open Tensor
open NN.Tensor
open NN.API

-- A scalar-valued function: mean(square x).
def sumsq : autograd.fn1.Fn (Shape.Vec 2) Shape.scalar :=
  fun x => do
    let y ← nn.functional.square x
    nn.functional.mean y

def demo : IO Unit := do
  let x : Tensor Float (Shape.Vec 2) := tensorND! [2] [0.5, -1.2]
  let g ← autograd.fn1.grad (α := Float) sumsq x
  IO.println s!"x    = {Spec.pretty x}"
  IO.println s!"grad = {Spec.pretty g}"

#check Proofs.Autograd.Algebra.backprop_correct
#check NN.Proofs.Autograd.Runtime.compileAux_ctx_eq_eval
#check NN.Proofs.Autograd.Runtime.backwardDenseFrom_compileAux_eq_backpropAllCtx

NN.Proofs.Autograd.Tape.Ops.Attention.ScaledDotProduct.backprop_eq_adjoint_fderiv_scaledDotProduct

NN.Proofs.Models.Attention.Weights.scaledDotProductAttention_weights_row_sum_one

3.3.9. Compiled Graphs🔗

The compiled runtime uses a typed SSA/DAG representation:

Proofs.Autograd.Algebra.GraphData

At this level, each node bundles:

• forward (compute a value),

• jvp (forward-mode pushforward),

• vjp (reverse-mode pullback).

These are precisely the ingredients needed to derive global correctness from local adjointness laws.

The reason to keep this compiled form separate from the eager tape is practical:

• eager mode is best for debugging and notebook-style iteration;

• compiled mode is best for a stable, replayable artifact that proof code can reason about;

• the same model definition can feed both paths, so the proof layer and the execution layer do not diverge.

TorchLean does not rely on a hidden global autograd engine. The current execution mode is part of the API, not an ambient effect.

opt = torch.optim.SGD(model.parameters(), lr=0.05)

-- conceptual shape, not exact syntax:
-- session owns parameters, inputs, tape, and backend choice
-- train.step applies one forward/backward/update step

3.3.10. Why Runtime and Verification Stay Separate🔗

The eager tape nodes are closures rather than symbolic op tags. That is excellent for execution, but much less suitable for verification passes such as IBP or CROWN, which require an explicit operator vocabulary.

TorchLean therefore keeps a first order IR with operation tags (NN.IR.Graph) and treats it as the canonical verification target. The compilation and reification path is covered in Graphs and IR.

3.3.11. Autograd Correctness Theorems🔗

TorchLean's autograd correctness is not a single monolithic theorem; it is layered into reusable pieces:

1. Local correctness of primitive ops The real autograd correctness API and semiring autograd correctness API state that each primitive VJP/JVP rule is the intended derivative or adjoint rule.

2. Global reverse mode correctness over a tape/DAG The tape algebra soundness API proves that composing locally correct nodes yields a globally correct reverse pass.

3. Link the executable runtime tape to the proved tape model The runtime autograd link API connects the executable tape representation to the proof model.

For a proof tour that parallels familiar PyTorch autograd concepts, start here:

• autograd proof overview

3.3.12. Debugging and Teaching🔗

The runtime layer is where widgets pay off:

• #tape_view and #tape_grads_view show recorded graphs and gradients.

• #tape_trace_view explains the reverse pass step by step (upstream cotangents and parent contributions).

• #runtime_ctx_view shows the value/gradient registries.

The Widgets guide and editor demo show these views in practice:

• widgets demo

3.3.13. Training Demos (Where To Look)🔗

TorchLean includes several PyTorch style training tutorials with selectable backends:

3.3.14. Optional CUDA / GPU execution🔗

Most of the semantic and verification layer in TorchLean is still pure Lean: specs, NN.IR.Graph, bound propagation, and theorems do not require a GPU. For training and large forward passes, the project can be linked against CUDA native code so float32 work runs on device buffers instead of only on host-side Float arrays.

This section is intentionally short. The full CUDA map is in GPU and CUDA.

lake build -R -K cuda=true

lake exe torchlean mlp --cuda --steps 20
lake exe torchlean gpt2 --cuda --steps 1
lake exe torchlean mamba --cuda --tiny-shakespeare --steps 25

3.3.15. Runtime Notes🔗

3.3.16. Common Pitfalls🔗

• Do not assume eager and compiled backends are interchangeable in every corner case. They aim to agree on supported programs, but they expose different debugging surfaces.

• Do not treat runtime tensors as proof objects. Runtime values are for execution and inspection; proof statements live in the spec/proof layers.

• Do not forget that heavy correctness proofs can be slow to elaborate. For everyday work, the executable checkers and widgets are usually the faster feedback loop.

To read the runtime layer in dependency order, start with eager tensors and tapes, then compiled graph construction, then the IR execution bridge, and finally the curated model and quickstart examples. That sequence usually preserves the runtime mental model better than jumping straight into the heaviest correctness proofs.

3.3.17. References🔗

• TorchLean paper (runtime + verification architecture): https://arxiv.org/abs/2602.22631

• Baydin et al., "Automatic differentiation in machine learning: a survey" (JMLR 2018): https://arxiv.org/abs/1502.05767

• Griewank and Walther, Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, the standard reference for local derivative rules composing into global reverse-mode correctness.

• PyTorch autograd docs (mental model + terminology): https://pytorch.org/docs/stable/autograd.html

• PyTorch autograd tutorial (gentle walkthrough): https://pytorch.org/tutorials/beginner/basics/autogradqs_tutorial.html

• PyTorch FX docs (graph capture mental model): https://pytorch.org/docs/stable/fx.html

• Runtime widgets: autograd widgets and runtime context widgets.