Shape and Concatenation IR Lowering

Checked lowering for permutations, reshaping, flattening, concatenation, and transpose.

Concatenation along axis 0 reads each parent through a typed index and folds Tensor.concatAxisSpec over the parents; concatenation along another axis first moves that axis to the front of every parent, folds, and moves it back. Both branches use concatLeadingAxisFromInputs, and the output-shape cast is justified by concatLeadingAxisFromInputs_size_eq_sum rather than by a proof embedded in the runtime code. The nonzero-axis branch validates each parent exactly as NN.IR.Graph.permuteSomeTensor does and records that evidence in ConcatFrontInput, so the correctness proof can replay the evaluator's permutation on every parent.

Each operation has its own small lower* definition. lowerShape only dispatches on the operation kind, and the lowerShape_* equation lemmas let correctness proofs reduce a dispatch to the branch they care about without unfolding the whole dispatcher.

def Runtime.Autograd.IRExec.Internal.lowerPermute {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (perm : Array ℕ) : NodeLoweringResult ctx

Checked lowering for .permute perm through a sequence of adjacent axis swaps.

def Runtime.Autograd.IRExec.Internal.lowerReshape {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (inS outS : Spec.Shape) : NodeLoweringResult ctx

Checked lowering for .reshape inS outS.

def Runtime.Autograd.IRExec.Internal.lowerFlatten {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (s : Spec.Shape) : NodeLoweringResult ctx

Checked lowering for .flatten s.

def Runtime.Autograd.IRExec.Internal.concatAxisZeroInputs {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (rest : Spec.Shape) : Except String (Array (ConcatInput α Γ rest))

Build the leading-axis concat inputs for .concat 0: one typed index per parent, all sharing the tail shape rest.

structure Runtime.Autograd.IRExec.Internal.ConcatFrontInput (α : Type) [TorchLean.Storage α] (Γ : List Spec.Shape) (permFront : Array ℕ) (restFront : Spec.Shape) : Type

One validated parent of a nonzero-axis concat: its typed index, the adjacent swaps that move the concatenated axis to the front, and the resulting leading extent. The proof fields record exactly the checks the IR evaluator performs in NN.IR.Graph.permuteSomeTensor, so the lowering and the evaluator permute each parent identically.

• sIn : Spec.Shape

  The parent's declared shape.

• ip : Proofs.Idx Γ self.sIn

  Typed index of the parent in the runtime context.

• swaps : Array ℕ

  Adjacent swaps realizing permFront on a tensor of rank sIn.rank.

• nP : ℕ

  Leading extent of the permuted parent.

• perm_eq : self.sIn.permute? permFront.toList = some (Spec.Shape.dim self.nP restFront)

  permFront is a valid permutation of the parent shape with the shared tail.

• swaps_eq : NN.IR.Graph.swapDepthsForPerm permFront self.sIn.rank = Except.ok self.swaps

  The swaps were computed for the parent's own rank.

• final_eq : swapShapeBySwaps self.sIn self.swaps = Spec.Shape.dim self.nP restFront

  Applying the swaps yields the permuted shape.

def Runtime.Autograd.IRExec.Internal.ConcatFrontInput.toInput {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} {permFront : Array ℕ} {restFront : Spec.Shape} (input : ConcatFrontInput α Γ permFront restFront) : ConcatInput α Γ restFront

Read and permute the parent so the concatenated axis comes first.

def Runtime.Autograd.IRExec.Internal.concatAxisFrontInputs {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (permFront : Array ℕ) (restFront : Spec.Shape) : Except String (Array (ConcatFrontInput α Γ permFront restFront))

Build the concat inputs for a nonzero axis: every parent is read through a typed index and permuted by permFront so the concatenated axis comes first, sharing the tail restFront.

def Runtime.Autograd.IRExec.Internal.concatInputsForward {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} {rest : Spec.Shape} (inputs : Array (ConcatInput α Γ rest)) (nOut : ℕ) (hSum : Array.foldl (fun (acc : ℕ) (input : ConcatInput α Γ rest) => acc + input.fst) 0 inputs = nOut) (context : TorchLean.TensorPack α Γ) : TorchLean.Tensor α (Spec.Shape.dim nOut rest)

Fold concat inputs along the leading axis into a tensor of the declared leading extent nOut.

hSum records that the input extents add up to nOut, so the cast is justified by concatLeadingAxisFromInputs_size_eq_sum.

def Runtime.Autograd.IRExec.Internal.lowerConcat {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (axis : ℕ) : NodeLoweringResult ctx

Checked lowering for .concat axis along an arbitrary axis.

def Runtime.Autograd.IRExec.Internal.lowerTranspose {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (axis₁ axis₂ : ℕ) : NodeLoweringResult ctx

Checked lowering for .transpose axis₁ axis₂ through a sequence of adjacent axis swaps.

def Runtime.Autograd.IRExec.Internal.lowerShape {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (kind : NN.IR.OpKind) : NodeLoweringResult ctx

Checked lowering for permutations, reshaping, flattening, concatenation, and transpose.

@[simp]

theorem Runtime.Autograd.IRExec.Internal.lowerShape_permute {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (perm : Array ℕ) : lowerShape ctx (NN.IR.OpKind.permute perm) = lowerPermute ctx perm

Dispatch equation for .permute perm.

@[simp]

theorem Runtime.Autograd.IRExec.Internal.lowerShape_reshape {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (inS outS : Spec.Shape) : lowerShape ctx (NN.IR.OpKind.reshape inS outS) = lowerReshape ctx inS outS

Dispatch equation for .reshape inS outS.

@[simp]

theorem Runtime.Autograd.IRExec.Internal.lowerShape_flatten {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (s : Spec.Shape) : lowerShape ctx (NN.IR.OpKind.flatten s) = lowerFlatten ctx s

Dispatch equation for .flatten s.

@[simp]

theorem Runtime.Autograd.IRExec.Internal.lowerShape_concat {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (axis : ℕ) : lowerShape ctx (NN.IR.OpKind.concat axis) = lowerConcat ctx axis

Dispatch equation for .concat axis.

@[simp]

theorem Runtime.Autograd.IRExec.Internal.lowerShape_transpose {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (axis₁ axis₂ : ℕ) : lowerShape ctx (NN.IR.OpKind.transpose axis₁ axis₂) = lowerTranspose ctx axis₁ axis₂

Dispatch equation for .transpose axis₁ axis₂.