LinkedSession

Proof-linked imperative Session (eager-style API, proved IR under the hood).

Background:

• Runtime.Autograd.TorchLean.Session provides a unified imperative API for training/debugging (eager) and verification-friendly execution (compiled).
• Proofs.Autograd.Algebra.GraphData is the proved/typed SSA(DAG) IR used by the proof-compiled pipeline (Proofs.Autograd.Algebra.Graph.compileAuxData), and NN/Proofs/Autograd/Runtime/Link.lean proves that running the runtime reverse-mode loop on the compiled tape matches GraphData.backpropAllCtx.

This file provides a session-style API that records a GraphData (well-typed IR) as you call ops imperatively, and then runs the standard runtime tape loop on the compiled tape.

Key guarantee (pure theorem, no IO reasoning needed):

• If the session snapshot is (g, x), then Tape.backwardDenseFrom (compileAuxData g x) equals GraphData.backpropAllCtx g x (via backwardDenseFrom_compileAuxData_eq_backpropAllCtx).

Practical note:

• This session enforces a simple invariant: all leaf tensors are created before any op node. This matches the standard training pattern (reset → add leaves → forward → backward).
• const is available as a graph node, so you can still introduce literal constants mid-graph.
• This is the fully proof-linked variant used by TorchLean.Session when opts.backend := .compiled.

@[reducible, inline]

abbrev Runtime.Autograd.Torch.Internal.okOrThrow {α : Type} : Result α → IO α

Convenience: turn a Result α into IO α by throwing IO.userError on .error.

This mirrors the common pattern in the eager runtime front-end (Torch.Core).

@[reducible, inline]

abbrev Runtime.Autograd.Torch.Internal.NatEnv : Type

Non-differentiable external environment for the proved graph: a small array of Nat inputs.

structure Runtime.Autograd.Torch.Internal.SessionIRState (α : Type) : Type

Internal proof-linked session state (a well-typed GraphData plus its leaf values).

• Γ : List Spec.Shape

  Leaf shapes (inputs/parameters), in creation order.

• x : Proofs.Autograd.Algebra.TList α self.Γ

  Leaf values, aligned with Γ.

• nat : NatEnv

  Non-differentiable external inputs (e.g. class labels/indices).

• ss : List Spec.Shape

  Internal node shapes, in creation order.

• g : Proofs.Autograd.Algebra.GraphData α NatEnv self.Γ self.ss

  SSA/DAG graph nodes (one per entry in ss).

def Runtime.Autograd.Torch.Internal.SessionIRState.empty {α : Type} : SessionIRState α

Empty session state: no leaves, no nodes, empty nat-environment.

structure Runtime.Autograd.Torch.Internal.SessionIR (α : Type) : Type

SessionIR is an imperative session that records a GraphData (proved IR) as it runs.

It is "eager-style" (you call ops imperatively), but it is proof-linked: the recorded graph can be compiled and then the runtime tape backward loop is provably equal to GraphData.backpropAllCtx.

• opts : Options

  Session options shared with the eager front-end.

• st : IO.Ref (SessionIRState α)

  Mutable proof-linked graph snapshot.

• paramsByLeaf : IO.Ref (Std.HashMap ℕ (AnyParam α))

  Map from graph leaf ids to mutable parameter objects.

def Runtime.Autograd.Torch.Internal.SessionIR.new {α : Type} (opts : Options := { }) : IO (SessionIR α)

Create a new proof-linked session.

This allocates IO.Refs for the session snapshot (SessionIRState) and the leaf-id→parameter map. Call resetTape to start a new "graph recording" phase.

def Runtime.Autograd.Torch.Internal.SessionIR.resetTape {α : Type} (s : SessionIR α) : IO Unit

Reset the session to an empty snapshot.

Important invariant: this session requires that all leaves are created before any op node. resetTape is the intended boundary between training steps/forwards.

def Runtime.Autograd.Torch.Internal.SessionIR.param {α : Type} (s : SessionIR α) {sh : Spec.Shape} (init : Spec.Tensor α sh) (name : Option String := none) (requiresGrad : Option Bool := none) : IO (Param α sh)

Create a mutable parameter object (not yet part of the recorded graph).

To use the parameter in the recorded graph, call use, which reads its current value and records it as a leaf in Γ. PyTorch comparison: analogous to creating a torch.nn.Parameter and then using it in a forward.

def Runtime.Autograd.Torch.Internal.SessionIR.ensureNoNodes {α : Type} (st : SessionIRState α) : IO Unit

Enforce the session invariant: leaves must be created before any op node.

This matches the usual training pattern: resetTape → add leaves → forward ops → backward.

def Runtime.Autograd.Torch.Internal.SessionIR.addLeaf {α : Type} (s : SessionIR α) {sh : Spec.Shape} (v : Spec.Tensor α sh) : IO (TensorRef α sh)

Record a new differentiable leaf tensor in the session context Γ.

This is the primitive used by use (parameters) and input (external inputs).

def Runtime.Autograd.Torch.Internal.SessionIR.use {α : Type} (s : SessionIR α) {sh : Spec.Shape} [DecidableEq Spec.Shape] (p : Param α sh) : IO (TensorRef α sh)

Use a Param in the recorded graph by reading its current value and recording it as a leaf.

The returned TensorRef is the graph handle you pass to subsequent ops. The session also remembers which leaf-id corresponds to which parameter, so sgdStepAll can update parameters after backward. PyTorch comparison: like referencing a torch.nn.Parameter in the forward; the parameter's value is treated as a leaf for autograd.

def Runtime.Autograd.Torch.Internal.SessionIR.input {α : Type} (s : SessionIR α) {sh : Spec.Shape} [DecidableEq Spec.Shape] (v : Spec.Tensor α sh) (name : Option String := none) (requiresGrad : Bool := false) : IO (TensorRef α sh)

Record an external differentiable input tensor as a leaf.

name and requiresGrad are accepted for API parity with the eager session, but this proof-linked session always records the input in Γ (a leaf) and uses typing/invariants to determine what gradients are meaningful.

def Runtime.Autograd.Torch.Internal.SessionIR.inputNat {α : Type} (s : SessionIR α) (v : ℕ) : IO NatRef

Record a non-differentiable Nat input in the external environment.

This is used for "index-like" inputs (labels, gather indices, etc.) that should not receive gradients. PyTorch comparison: like passing an integer tensor / index to an op; indices are not differentiable.

def Runtime.Autograd.Torch.Internal.SessionIR.getNat {α : Type} (s : SessionIR α) (r : NatRef) : IO ℕ

Read a previously recorded NatRef.

def Runtime.Autograd.Torch.Internal.SessionIR.setNat {α : Type} (s : SessionIR α) (r : NatRef) (v : ℕ) : IO Unit

Overwrite a previously recorded NatRef.

def Runtime.Autograd.Torch.Internal.SessionIR.natVecToArray {k : ℕ} (v : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) : Array ℕ

Convert a small Tensor Nat (.dim k .scalar) into an Array Nat.

This is used to stage NatVecRef inputs into the session nat-environment.

def Runtime.Autograd.Torch.Internal.SessionIR.inputNatVec {α : Type} {k : ℕ} (s : SessionIR α) (v : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) : IO (NatVecRef k)

Record a non-differentiable vector of Nat inputs.

Returns a NatVecRef k which points into the nat-environment. This is useful for "runtime gather" style ops where indices are supplied externally (and are not differentiable).

def Runtime.Autograd.Torch.Internal.SessionIR.getNatVec {α : Type} {k : ℕ} (s : SessionIR α) (r : NatVecRef k) : IO (Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar))

Read back the k-vector stored at a NatVecRef k.

def Runtime.Autograd.Torch.Internal.SessionIR.setNatVec {α : Type} {k : ℕ} (s : SessionIR α) (r : NatVecRef k) (v : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) : IO Unit

Overwrite the nat-environment segment referenced by NatVecRef k.

def Runtime.Autograd.Torch.Internal.SessionIR.mkIdxOrThrow {_α : Type} {Γ ss : List Spec.Shape} (id : ℕ) (s : Spec.Shape) : Result (Proofs.Autograd.Algebra.Idx (Γ ++ ss) s)

Build a typed index into the current context Γ ++ ss from a raw numeric id and expected shape.

This is the main "dynamic check" used by getValue (and by a few index-driven nodes): it ensures that the Nat id points to an existing tensor in the session context and that the shape matches.

def Runtime.Autograd.Torch.Internal.SessionIR.getValue {α : Type} (s : SessionIR α) {sh : Spec.Shape} [DecidableEq Spec.Shape] (x : TensorRef α sh) : IO (Spec.Tensor α sh)

Evaluate the recorded graph and return the value of a TensorRef.

This is a pure graph evaluation (GraphData.eval) using the recorded leaf values and nat-environment. It does not run the runtime tape or mutate session state.