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.CompiledScalar (α : 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 α Unit Γ self.ssPrev

  Proved graph prefix that computes all preceding SSA nodes.

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

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

@[reducible, inline]

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

Convenience alias for the proved heterogeneous tensor list over a shape context.

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.CompiledOut (α : 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.

@[reducible, inline]

abbrev Runtime.Autograd.Torch.CompiledOut.TList (α : Type) (ss : List Spec.Shape) : Type

Convenience alias for the proved heterogeneous tensor list over a shape context.

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

Evaluate the output tensor for leaf values x.

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

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

def Runtime.Autograd.Torch.CompiledOut.vjpWithSeed {α : Type} [Add α] [Zero α] {Γ : List Spec.Shape} {τ : Spec.Shape} (c : CompiledOut α Γ τ) (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.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.compileOut {α : Type} [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {τ : Spec.Shape} (build : Compiled.GraphM.M α Γ (Compiled.GraphM.Var τ)) : Result (CompiledOut α Γ τ)

Compile a tensor-output graph builder into a CompiledOut.

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.