Proof-Linked Session: Differentiation and Backpropagation

Backward + SGD (runtime tape loop on the compiled tape)

def Runtime.Autograd.Torch.Internal.SessionIR.compileTape {α : Type} [DecidableEq Spec.Shape] (st : SessionIRState α) : Tape α

Compile the recorded proved graph into a runtime tape.

This uses Graph.compileAuxData (the same compiler used by the proof pipeline) and extracts the runtime tape component.

def Runtime.Autograd.Torch.Internal.SessionIR.backwardDenseAll {α : Type} (s : SessionIR α) [Add α] [Zero α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (out : TensorRef α sh) (seed : Spec.Tensor α sh) : IO (Array (AnyTensor α))

Run reverse-mode backprop for the whole recorded context and return a dense gradient array.

seed is the upstream gradient for out (same convention as PyTorch's loss.backward(gradient=...)).

def Runtime.Autograd.Torch.Internal.SessionIR.backwardScalarDenseAll {α : Type} (s : SessionIR α) [Add α] [Zero α] [One α] [DecidableEq Spec.Shape] (loss : TensorRef α Spec.Shape.scalar) : IO (Array (AnyTensor α))

Run backward from a scalar loss with seed 1.

PyTorch comparison: loss.backward() for a scalar loss.

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

Extract the gradient tensor for a particular TensorRef from a dense gradient array.

This is the typed analogue of looking up grads[x.id] and casting it to the expected shape.

Forward-mode: JVP (compiled only)

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

Like mkIdxOrThrow, but restricted to leaves Γ only.

def Runtime.Autograd.Torch.Internal.SessionIR.dxTListFromAnyArray {α : Type} [Zero α] [DecidableEq Spec.Shape] (Γ : List Spec.Shape) (dxs : Array (AnyTensor α)) : IO (Proofs.Autograd.Algebra.TList α Γ)

Convert a dense tangent array (aligned with leaf creation order) into a typed TList α Γ.

This is the main adapter needed to call the proved GraphData.jvpCtx forward-mode routine.

def Runtime.Autograd.Torch.Internal.SessionIR.dxTListFromAnyArray.go {α : Type} [DecidableEq Spec.Shape] (dxs : Array (AnyTensor α)) (Γ' : List Spec.Shape) (off : ℕ) : IO (Proofs.Autograd.Algebra.TList α Γ')

def Runtime.Autograd.Torch.Internal.SessionIR.jvpDenseAll {α : Type} (s : SessionIR α) [Zero α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (out : TensorRef α sh) (dxs : Array (AnyTensor α)) : IO (Spec.Tensor α sh)

Jacobian-vector product for the current session snapshot.

dxs is a dense array of tangents for leaf tensors, aligned with leaf creation order.

def Runtime.Autograd.Torch.Internal.SessionIR.jvpLeaf {α : Type} (s : SessionIR α) [Zero α] [DecidableEq Spec.Shape] {shOut shX : Spec.Shape} (out : TensorRef α shOut) (x : TensorRef α shX) (dx : Spec.Tensor α shX) : IO (Spec.Tensor α shOut)

JVP for a single leaf: tangent is nonzero only at x.

def Runtime.Autograd.Torch.Internal.SessionIR.jvpScalarLeaf {α : Type} (s : SessionIR α) [Zero α] [DecidableEq Spec.Shape] (loss : TensorRef α Spec.Shape.scalar) {shX : Spec.Shape} (x : TensorRef α shX) (dx : Spec.Tensor α shX) : IO α

Scalar-loss JVP for a single leaf.

def Runtime.Autograd.Torch.Internal.SessionIR.sgdStepAll {α : Type} (s : SessionIR α) [Sub α] [Mul α] [Add α] [Zero α] [DecidableEq Spec.Shape] (lr : α) (grads : Array (AnyTensor α)) : IO Unit

Apply an SGD update to all parameters recorded via use.

grads is expected to be the dense gradient array returned by backwardDenseAll / backwardScalarDenseAll. Only entries corresponding to parameters (leaves that were produced by use) are used to update Param.value. PyTorch comparison: like iterating params and doing p.data -= lr * p.grad.

Pure correctness hook: session snapshot ↔ proved IR backprop

theorem Runtime.Autograd.Torch.Internal.SessionIR.backwardDenseFrom_compileAuxData_eq_backpropAllCtx {α : Type} [DecidableEq Spec.Shape] [CommSemiring α] (st : SessionIRState α) (seed : Proofs.Autograd.Algebra.TList α (st.Γ ++ st.ss)) : (Proofs.Autograd.Algebra.Graph.compileAuxData st.g st.x st.nat).1.backwardDenseFrom seed.toAnyArray = Except.ok (st.g.backpropAllCtx st.x st.nat seed).toAnyArray

Core proof-link: running the runtime reverse-mode loop on the compiled tape equals proved backprop.

This theorem is the "hook" that lets a session-style API be backed by the proved IR: compileAuxData produces a tape, and Tape.backwardDenseFrom is shown equal to GraphData.backpropAllCtx (up to the TList.toAnyArray representation change).