Compiled Torch Wrappers

Thin torch.compile-style wrappers around the proof-backed graph compiler. These definitions expose forward, JVP, and VJP/backward entry points while keeping the compiled graph semantics explicit.

Imperative sessions live in:

• Runtime.Autograd.TorchLean.Session (unified eager/compiled, recommended),
• Runtime.Autograd.Torch.Internal.SessionIR (proof-linked imperative session, internal).

torch.compile-style wrapper (cached static graph).

This record contains the proof-compiled graph model (GraphData) and its proven-correct reverse-mode accumulator (GraphData.backpropCtx). It compiles once, then you can call forward/backward repeatedly with new inputs.

Note: this does not cache a Runtime.Autograd.Tape for reuse across different inputs. The current tape compiler bakes the forward context into backward closures, so reusing a single tape across changing inputs would be unsound without redesigning the runtime node API.

structure Runtime.Autograd.Torch.CompiledScalarWith (α Δ : Type) (Γ : List Spec.Shape) : Type

torch.compile-style wrapper for a scalar-valued computation over leaf context Γ.

This stores a proved node (NodeData) together with the preceding graph prefix so it can be evaluated and differentiated without rebuilding the whole graph each time.

• ssPrev : List Spec.Shape

  Shapes of internal SSA nodes preceding the scalar output node.

• gPrev : Proofs.Autograd.Algebra.GraphData α Δ Γ self.ssPrev

  Proved graph prefix that computes all preceding SSA nodes.

• node : Proofs.Autograd.Algebra.NodeData α Δ (Γ ++ self.ssPrev) Spec.Shape.scalar

  Final scalar output node over the leaf context plus graph prefix.

@[reducible, inline]

abbrev Runtime.Autograd.Torch.CompiledScalar (α : Type) (Γ : List Spec.Shape) : Type

Scalar compiled graph with no auxiliary, non-differentiable runtime inputs.

def Runtime.Autograd.Torch.CompiledScalar.forward {α : Type} {Γ : List Spec.Shape} (c : CompiledScalar α Γ) (x : TList α Γ) : Spec.Tensor α Spec.Shape.scalar

Evaluate the scalar output for leaf values x.

def Runtime.Autograd.Torch.CompiledScalar.jvp {α : Type} {Γ : List Spec.Shape} (c : CompiledScalar α Γ) (x dx : TList α Γ) : Spec.Tensor α Spec.Shape.scalar

Forward-mode Jacobian-vector product (JVP) at x with tangent dx.

def Runtime.Autograd.Torch.CompiledScalar.backward {α : Type} [Add α] [Zero α] [One α] {Γ : List Spec.Shape} (c : CompiledScalar α Γ) (x : TList α Γ) : TList α Γ

Reverse-mode backprop for a scalar output with implicit seed 1.

Returns a TList of gradients aligned with the leaf context Γ.

def Runtime.Autograd.Torch.CompiledScalar.backwardWithSeed {α : Type} [Add α] [Zero α] {Γ : List Spec.Shape} (c : CompiledScalar α Γ) (x : TList α Γ) (seedOut : α) : TList α Γ

Reverse-mode backprop for a scalar output with an explicit scalar seed.

PyTorch comparison: loss.backward(gradient=seedOut) for a scalar loss.

torch.compile-style wrapper for tensor-valued outputs.

This is the same idea as CompiledScalar, but parameterized by an arbitrary output shape τ. It supports:

• forward (evaluate the output),
• jvp (forward-mode JVP, provided all ops supply jvp),
• vjpWithSeed (reverse-mode VJP with an explicit cotangent seed at the output).

structure Runtime.Autograd.Torch.CompiledGraph (α : Type) (Γ : List Spec.Shape) (τ : Spec.Shape) : Type

torch.compile-style wrapper for a tensor-valued output of shape τ.

This generalizes CompiledScalar to arbitrary output shapes and provides forward-mode JVP and reverse-mode VJP (with explicit seed).

• ssPrev : List Spec.Shape

  Shapes of internal SSA nodes preceding the output node.

• gPrev : Proofs.Autograd.Algebra.GraphData α Unit Γ self.ssPrev

  Proved graph prefix that computes all preceding SSA nodes.

• node : Proofs.Autograd.Algebra.NodeData α Unit (Γ ++ self.ssPrev) τ

  Final output node over the leaf context plus graph prefix.

def Runtime.Autograd.Torch.CompiledGraph.forward {α : Type} {Γ : List Spec.Shape} {τ : Spec.Shape} (c : CompiledGraph α Γ τ) (x : TList α Γ) : Spec.Tensor α τ

Evaluate the output tensor for leaf values x.

def Runtime.Autograd.Torch.CompiledGraph.jvp {α : Type} {Γ : List Spec.Shape} {τ : Spec.Shape} (c : CompiledGraph α Γ τ) (x dx : TList α Γ) : Spec.Tensor α τ

Forward-mode Jacobian-vector product (JVP) at x with tangent dx.

def Runtime.Autograd.Torch.CompiledGraph.vjpWithSeed {α : Type} [Add α] [Zero α] {Γ : List Spec.Shape} {τ : Spec.Shape} (c : CompiledGraph α Γ τ) (x : TList α Γ) (seedOut : Spec.Tensor α τ) : TList α Γ

Reverse-mode vector-Jacobian product (VJP) with an explicit output cotangent seed.

This is the tensor-valued analogue of CompiledScalar.backwardWithSeed. PyTorch comparison: out.backward(gradient=seedOut) (for a tensor output).

def Runtime.Autograd.Torch.compileScalarWith {α Δ : Type} [DecidableEq Spec.Shape] {Γ : List Spec.Shape} (build : Compiled.GraphM.MWith α Δ Γ (Compiled.GraphM.Var Spec.Shape.scalar)) : Result (CompiledScalarWith α Δ Γ)

Compile a scalar-output graph whose nodes may read an auxiliary runtime environment Δ.

The environment is not part of the differentiable leaf context Γ. It is intended for discrete inputs such as token ids, labels, gather indices, and masks that affect evaluation but do not receive gradients.

def Runtime.Autograd.Torch.compileScalarWith.finish {α Δ : Type} {Γ : List Spec.Shape} (outVar : Compiled.GraphM.Var Spec.Shape.scalar) (ss : List Spec.Shape) : Compiled.GraphM.PGraphData α Δ Γ ss → Result (CompiledScalarWith α Δ Γ)

def Runtime.Autograd.Torch.compileScalar {α : Type} [DecidableEq Spec.Shape] {Γ : List Spec.Shape} (build : Compiled.GraphM.M α Γ (Compiled.GraphM.Var Spec.Shape.scalar)) : Result (CompiledScalar α Γ)

Compile a scalar-output graph builder into a CompiledScalar.

The builder is expressed in the Compiled.GraphM monad. We expect it to produce at least one node and return a variable of scalar shape.

def Runtime.Autograd.Torch.compileGraph {α : Type} [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {τ : Spec.Shape} (build : Compiled.GraphM.M α Γ (Compiled.GraphM.Var τ)) : Result (CompiledGraph α Γ τ)

Compile a tensor-output graph builder into a CompiledGraph.

We require that the returned Var τ is the last node produced by the builder, so the wrapper can store the prefix graph and final output node cleanly.

def Runtime.Autograd.Torch.compileGraph.finish [DecidableEq Spec.Shape] {α : Type} {Γ : List Spec.Shape} {τ : Spec.Shape} (outVar : Compiled.GraphM.Var τ) (ss : List Spec.Shape) : Compiled.GraphM.PGraphData α Unit Γ ss → Result (CompiledGraph α Γ τ)