IR Graph Lowering

Checked lowering from the shared operation-tagged IR into executable proof-compiled graph data. The representation and denotation helpers live in IRExec.Core.

@[reducible, inline]

abbrev Runtime.Autograd.Compiled.IRExec.State (α : Type) (inShape : Spec.Shape) : Type

Internal compilation state used by buildFrom.

It is a dependent pair of:

• ss: shapes of already-compiled IR nodes,
• GraphData α Unit [inShape] ss: executable closures for exactly that shape list.

def Runtime.Autograd.Compiled.IRExec.mkIdx [DecidableEq Spec.Shape] (inShape : Spec.Shape) (ss : List Spec.Shape) (id : ℕ) (s : Spec.Shape) : Except String (Proofs.Autograd.Algebra.Idx ([inShape] ++ ss) s)

Build a typed runtime index (Idx) for a numeric IR parent id.

The compiled runtime context is typed by [inShape] ++ ss, matching GraphData's input-plus-node context representation. mkIdx checks that:

• id is in bounds, and
• the context shape at that position matches the expected shape s.

On failure, this returns a descriptive error string used directly by buildFrom.

def Runtime.Autograd.Compiled.IRExec.getIRValue {α : Type} {inShape : Spec.Shape} {ss : List Spec.Shape} {s : Spec.Shape} (ctx : Proofs.Autograd.Algebra.TList α ([inShape] ++ ss)) (idx : Proofs.Autograd.Algebra.Idx ([inShape] ++ ss) s) : Spec.Tensor α s

Read a tensor from the single-input IR execution context using a checked parent index.

Keeping the context spelling [inShape] ++ ss explicit prevents dependent elaboration from normalizing the singleton append differently at compiler and correctness-proof call sites.

def Runtime.Autograd.Compiled.IRExec.mkFwdNode {α : Type} [Context α] {Γ : List Spec.Shape} {τ : Spec.Shape} (forward : Proofs.Autograd.Algebra.TList α Γ → Spec.Tensor α τ) : Proofs.Autograd.Algebra.NodeData α Unit Γ τ

Construct a NodeData for forward execution only.

The compiled runtime GraphData expects each node to supply forward, jvp, and vjp. For this IR bridge we only care about forward correctness, so jvp/vjp are populated with forward-only sentinels that panic! if called.

This is intentional: IRExec closes the forward semantics gap; full gradient behavior is handled by other runtime/autograd layers. Using panic! here is a safety measure: it prevents silently wrong gradients if someone accidentally routes differentiation through an IRExec-compiled graph.

def Runtime.Autograd.Compiled.IRExec.mkFwdNode.defaultTList {α : Type} [Context α] {Γ : List Spec.Shape} : Proofs.Autograd.Algebra.TList α Γ

@[simp]

theorem Runtime.Autograd.Compiled.IRExec.mkFwdNode_forward {α : Type} [Context α] {Γ : List Spec.Shape} {τ : Spec.Shape} (f : Proofs.Autograd.Algebra.TList α Γ → Spec.Tensor α τ) (ctx : Proofs.Autograd.Algebra.TList α Γ) (d : Unit) : (mkFwdNode f).forward ctx d = f ctx

Forward projection for mkFwdNode.

The JVP/VJP fields are sentinels in this bridge, but the forward field is exactly the function passed to the constructor. This small simp lemma is used by the IR semantic-equivalence proof.

def Runtime.Autograd.Compiled.IRExec.swapShapeBySwaps (s : Spec.Shape) : List ℕ → Spec.Shape

Apply a list of adjacent swaps (specified by swap depths) to a shape.

This is the shape-level companion of applySwapsTensor, and mirrors IR permutation lowering.

def Runtime.Autograd.Compiled.IRExec.applySwapsTensor {α : Type} [Context α] {s : Spec.Shape} (swaps : List ℕ) : Spec.Tensor α s → Spec.Tensor α (swapShapeBySwaps s swaps)

Apply the same swap sequence as swapShapeBySwaps, but to a tensor value.

This uses Tensor.swap_at_depth_helper repeatedly; it is the runtime companion of the IR-side swapDepthsForPerm lowering used by .permute.

def Runtime.Autograd.Compiled.IRExec.concatLeadingAxisFromInfos {α : Type} [Context α] {Γ : List Spec.Shape} {rest : Spec.Shape} (ctx : Proofs.Autograd.Algebra.TList α Γ) (infos : List ((nP : ℕ) × Proofs.Autograd.Algebra.Idx Γ (Spec.Shape.dim nP rest))) : (nSum : ℕ) × Spec.Tensor α (Spec.Shape.dim nSum rest)

Concatenate a list of tensors (all with shape .dim nP rest) along dimension 0.

The input list is expressed as typed indices into the runtime context Γ; the result tracks the total concatenated size as a sigma.

This helper supports lowering of IR concat-style operators while preserving shape information.

theorem Runtime.Autograd.Compiled.IRExec.concatLeadingAxisFromInfos_size_eq_sum {α : Type} [Context α] {Γ : List Spec.Shape} {rest : Spec.Shape} (ctx : Proofs.Autograd.Algebra.TList α Γ) (infos : List ((nP : ℕ) × Proofs.Autograd.Algebra.Idx Γ (Spec.Shape.dim nP rest))) : (concatLeadingAxisFromInfos ctx infos).fst = List.foldl (fun (acc : ℕ) (info : (nP : ℕ) × Proofs.Autograd.Algebra.Idx Γ (Spec.Shape.dim nP rest)) => acc + info.fst) 0 infos

The concatenated size reported by concatLeadingAxisFromInfos is the sum of the input sizes.

This theorem is used to justify the output-shape side conditions in concat lowering branches.