Certified einsum operand indexing

This module generates row-major operand indices, hoists reusable coordinate contributions, and certifies direct tensor reads against the checked einsum index plan.

def TorchLean.Tensor.Internal.Elab.Impl.shouldHoistOutputIndexTerms (inputAxes : List Check.EinsumAxis) (inputDimensions : List Lean.Expr) (outputAxes : List Check.EinsumAxis) : Lean.MetaM Bool

Hoist multi-term output-coordinate stride products out of contraction loops.

Lean's native compiler does not reliably perform loop-invariant motion for boxed natural-number arithmetic. Precomputing a multi-term base therefore removes several products and additions from every contracted iteration. Passing a one-term base through the generated loop is measurably slower than recomputing that product, so singleton dimensions and strides that reduce to zero or one are ignored and at least two remaining terms are required.

def TorchLean.Tensor.Internal.Elab.Impl.addInputIndexContribution (physicalLength logicalLength base contribution : Lean.Expr) : Lean.MetaM Lean.Expr

Add one physical-axis contribution to a row-major input index.

A literal singleton contributes nothing. When the physical length is definitionally the logical axis length, the coordinate already belongs to that same Fin type and is necessarily zero in the singleton case, so no runtime broadcast branch is needed. Only genuinely ambiguous symbolic broadcasts retain the equality test.

def TorchLean.Tensor.Internal.Elab.Impl.hoistedOutputIndexBase (inputDimensions logicalDimensions : List Lean.Expr) (outputCoordinates : List (Option Lean.Expr)) : Lean.MetaM Lean.Expr

Combine one operand's loop-invariant output contributions in row-major index order.

Unknown physical dimensions retain the singleton-broadcast guard outside the contraction loop. The generated index certificate treats natural-number addition as associative and commutative, so output contributions can be regrouped without changing the verified flat index.

def TorchLean.Tensor.Internal.Elab.Impl.contractedAxisIndexContribution (contractedAxis : Check.EinsumAxis) (contractedCoordinate : Lean.Expr) (inputAxes : List Check.EinsumAxis) (inputDimensions logicalDimensions : List Lean.Expr) : Lean.MetaM Lean.Expr

Collect one contracted logical axis's physical row-major contribution to an operand index.

Repeated labels contribute once per physical occurrence. Symbolic singleton dimensions retain the runtime broadcast guard required by the verified plan.

def TorchLean.Tensor.Internal.Elab.Impl.compileInputFlatIndexValue (inputAxes : List Check.EinsumAxis) (inputDimensions logicalDimensions : List Lean.Expr) (outputCoordinates : List (Option Lean.Expr)) (hoistedOutputBase : Option Lean.Expr) (hoistedContractedAxes : List Check.EinsumAxis) (contractionCoordinates : List (Check.EinsumAxis × Lean.Expr)) : Lean.MetaM Lean.Expr

Generate one operand's physical row-major index, certified against the generic index plan. Natural-number index sums may be regrouped; this does not reorder the scalar contraction.

For operands selected by shouldHoistOutputIndexTerms, all loop-invariant output contributions arrive as one precomputed base. Contracted axes already included in a staged base are omitted from the remaining dynamic index.

def TorchLean.Tensor.Internal.Elab.Impl.finIndexEquality (inputSize inputIndex certifiedInputIndex hIndexValue : Lean.Expr) : Lean.MetaM Lean.Expr

Lift equality of two bounded index values to equality of their Fin terms.

def TorchLean.Tensor.Internal.Elab.Impl.compileInputRead (tensor inputSize inputIndexValue normalizedInputIndexValue certifiedInputIndex hInputIndexValue : Lean.Expr) (useNativeIndex : Bool) (directIndexBound? : Option Lean.Expr := none) (nativeIndexAssumptions : Array Lean.Expr := #[]) : Lean.Elab.TermElabM (Lean.Expr × Lean.Expr)

Compile one operand read from a generated row-major index.

Both branches end at the same certified Fin index. Concrete buffers that fit every Lean target use native word arithmetic and Storage.uget; buffers whose sizes remain symbolic or are at least 2^32 retain ordinary natural-number arithmetic. The native branch proves that USize.toNat recovers the generated reference index exactly, so wrapping cannot affect execution.