IR Lowering Primitives

Typed state, indices, and shape operations used by the checked IR lowering pass. These declarations remain under IRExec.Internal; the public lowering entry point lives in IRExec.Lowering.

@[reducible, inline]

abbrev Runtime.Autograd.IRExec.Internal.State (α : Type) [TorchLean.Storage α] (inShape : Spec.Shape) : Type

Internal lowering state used by buildFrom.

It is a dependent pair of:

• ss: shapes of already-lowered IR nodes,
• ForwardData α [inShape] ss: forward closures for exactly that shape list.

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

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

The forward executor's context is typed by [inShape] ++ ss, matching ForwardData's input-plus-node 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.IRExec.Internal.mkForwardNode {α : Type} [TorchLean.Storage α] {Γ : List Spec.Shape} {τ : Spec.Shape} (forward : TorchLean.TensorPack α Γ → TorchLean.Tensor α τ) : ForwardNode α Γ τ

Package a typed forward closure as one node of the executable IR graph.

@[simp]

theorem Runtime.Autograd.IRExec.Internal.mkForwardNode_eval {α : Type} [TorchLean.Storage α] {Γ : List Spec.Shape} {τ : Spec.Shape} (f : TorchLean.TensorPack α Γ → TorchLean.Tensor α τ) (ctx : TorchLean.TensorPack α Γ) : (mkForwardNode f).eval ctx = f ctx

Evaluation projection for mkForwardNode.

def Runtime.Autograd.IRExec.Internal.swapShapeBySwapsList (s : Spec.Shape) : List ℕ → Spec.Shape

Internal list recursion used to track the dependent output shape of adjacent swaps.

def Runtime.Autograd.IRExec.Internal.swapShapeBySwaps (s : Spec.Shape) (swaps : Array ℕ) : Spec.Shape

Apply adjacent swaps, represented by their axis depths, to a shape.

def Runtime.Autograd.IRExec.Internal.applySwapsTensorList {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (swaps : List ℕ) : TorchLean.Tensor α s → TorchLean.Tensor α (swapShapeBySwapsList s swaps)

Internal dependent recursion underlying applySwapsTensor.

def Runtime.Autograd.IRExec.Internal.applySwapsTensor {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (swaps : Array ℕ) (tensor : TorchLean.Tensor α s) : TorchLean.Tensor α (swapShapeBySwaps s swaps)

Apply the same adjacent-swap sequence as swapShapeBySwaps to a tensor value.

@[reducible, inline]

abbrev Runtime.Autograd.IRExec.Internal.ConcatInput (α : Type) [TorchLean.Storage α] (Γ : List Spec.Shape) (rest : Spec.Shape) : Type

One typed concat input: a leading extent together with a closure that reads the tensor with that extent from the runtime context. Inputs for a nonzero concat axis permute the parent before returning it, so the closure is the common shape for every concat branch.

def Runtime.Autograd.IRExec.Internal.concatLeadingAxisList {α : Type} [TorchLean.Storage α] [Context α] {rest : Spec.Shape} : List ((n : ℕ) × TorchLean.Tensor α (Spec.Shape.dim n rest)) → (nSum : ℕ) × TorchLean.Tensor α (Spec.Shape.dim nSum rest)

Concatenate typed tensors along their leading axis, folding from the first tensor.

The empty list yields the empty tensor with leading extent 0. This is the same fold shape as the IR evaluator's NN.IR.Graph.evalConcatLeadingAxisFold.

theorem Runtime.Autograd.IRExec.Internal.concatLeadingAxisList_fst {α : Type} [TorchLean.Storage α] [Context α] {rest : Spec.Shape} (tensors : List ((n : ℕ) × TorchLean.Tensor α (Spec.Shape.dim n rest))) : (concatLeadingAxisList tensors).fst = List.foldl (fun (acc : ℕ) (t : (n : ℕ) × TorchLean.Tensor α (Spec.Shape.dim n rest)) => acc + t.fst) 0 tensors

The leading extent of the fold in concatLeadingAxisList is a plain sum of extents.

def Runtime.Autograd.IRExec.Internal.concatLeadingAxisFromInputs {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} {rest : Spec.Shape} (ctx : TorchLean.TensorPack α Γ) (inputs : Array (ConcatInput α Γ rest)) : (nSum : ℕ) × TorchLean.Tensor α (Spec.Shape.dim nSum rest)

Concatenate the tensors produced by concat inputs along their leading axis.

theorem Runtime.Autograd.IRExec.Internal.concatLeadingAxisFromInputs_size_eq_sum {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} {rest : Spec.Shape} (ctx : TorchLean.TensorPack α Γ) (inputs : Array (ConcatInput α Γ rest)) : (concatLeadingAxisFromInputs ctx inputs).fst = Array.foldl (fun (acc : ℕ) (input : ConcatInput α Γ rest) => acc + input.fst) 0 inputs

The concatenated size reported by concatLeadingAxisFromInputs is the sum of the input extents.

This theorem justifies the output-shape cast in the concat lowering branches.