Fused native lowering for einsum

A checked einsum allocates only its final output tensor. For each output entry, the executable kernel:

1. enumerates flat indices for only the contracted logical axes;
2. evaluates a precompiled row-major stride plan for each operand;
3. reads every operand directly from its native array, with diagonal selection and singleton broadcasting already encoded by that plan;
4. multiplies the operand values in source order; and
5. sums those products into the output entry.

No broadcast, pointwise-product, permutation, or contraction tensor is materialized. The independent product tensor lives in the semantics module; the lowering module contains only the executable flat-index plan and its correctness bridge.

The executable loops need only addition, multiplication, and the scalar literals zero and one. Their row-major order is part of the program, which makes the same kernel available to IEEE floating-point types without installing false algebraic instances. Stronger correctness theorems recover the order-independent Rep.push denotation whenever the scalar operations form additive and multiplicative monoids.

References

The broadcast, ordered product, axis permutation, and reduction stages follow the contraction behavior of einops v0.8.2 at commit 8e911db71f2e693a0c434b041180388c685ed06f.

theorem TorchLean.Tensor.Internal.Lowering.contractedAxes_nodup (checked : Check.CheckedEinsum) : checked.contractedAxes.Nodup

Contracted axes inherit duplicate-freedom from the checked global axis list.

Verified flat indexing

The fused kernel compiles each operand's named axes to a short list of ordinary triples:

• whether the coordinate comes from the output or contraction coordinate;
• the position of that logical axis in its source coordinate; and
• the operand's physical row-major stride.

Singleton physical dimensions contribute no term. Evaluation then performs only coordinate projection, multiplication, addition, and one direct array read. The helper theorems below prove that this arithmetic index is exactly the linearization of inputCoordinateOfGlobal; all shape and broadcasting proofs are erased from generated code.

def TorchLean.Tensor.Internal.Lowering.coordinateAt (shape : Shape) : Coord shape → ℕ → ℕ

Read one component of a shape coordinate by its zero-based axis position.

def TorchLean.Tensor.Internal.Lowering.inputFlatIndexPlan {ι : Type u_1} [BEq ι] [LawfulBEq ι] (outputAxes contractedAxes inputAxes : List ι) : Shape → List (Bool × ℕ × ℕ)

Compile physical operand axes to source-position and input-stride triples.

The plan is built once per operand. Omitting singleton dimensions makes broadcasting free inside the scalar loop.

def TorchLean.Tensor.Internal.Lowering.evaluateInputFlatIndexPlan (outputShape contractedShape : Shape) (outputCoordinate : Coord outputShape) (contractionCoordinate : Coord contractedShape) : List (Bool × ℕ × ℕ) → ℕ

Evaluate one compiled operand plan to its physical row-major array index.

@[inline]

def TorchLean.Tensor.Internal.Lowering.einsumInputFlatIndex (checked : Check.CheckedEinsum) (operand : Fin checked.inputShapes.length) (outputCoordinate : Coord checked.output) (contractionCoordinate : Coord (List.map checked.axisLength checked.contractedAxes)) : Fin (checked.inputShapes.get operand).size

The certified physical row-major index read from one einsum operand.

This is the indexing contract shared by the generic kernel and pattern-specialized elaboration. Its value is ordinary natural-number arithmetic; the bound proof is erased from executable code.

The executable kernel uses nested Fin.foldl loops so contraction and operand traversal allocate no temporary finite sets or lists. The row-major traversal itself is Semantics.coordinateSum: there is one such loop in the library, and the lemmas below are what connect it to the big-operator semantics.

theorem TorchLean.Tensor.Internal.Lowering.coordinateSum_congr {R : Type u} [Add R] [OfNat R 0] (shape : Shape) (left right : Coord shape → R) (initial : R) (h : ∀ (coordinate : Coord shape), left coordinate = right coordinate) : Semantics.coordinateSum shape left initial = Semantics.coordinateSum shape right initial

Pointwise-equal scalar kernels have equal multidimensional coordinate sums.

theorem TorchLean.Tensor.Internal.Lowering.coordinateSum_congr_of_isEmpty {R : Type u} [Add R] [OfNat R 0] (shape : Shape) [IsEmpty (Coord shape)] (left right : Coord shape → R) (initial : R) : Semantics.coordinateSum shape left initial = Semantics.coordinateSum shape right initial

Every scalar kernel has the same coordinate sum over an empty coordinate space. This is the certificate used when an einsum contracts a zero-length axis, where the scalar kernels are intentionally never evaluated.

theorem TorchLean.Tensor.Internal.Lowering.foldl_mul_split_three {R : Type u} [Monoid R] (left middle right : List R) : List.foldl (fun (x1 x2 : R) => x1 * x2) 1 left * List.foldl (fun (x1 x2 : R) => x1 * x2) 1 middle * List.foldl (fun (x1 x2 : R) => x1 * x2) 1 right = List.foldl (fun (x1 x2 : R) => x1 * x2) 1 (left ++ middle ++ right)

Split an ordered multiplicative fold into three contiguous operand ranges.

The factors retain their source order; only associativity and the identity laws of a monoid are used.

theorem TorchLean.Tensor.Internal.Lowering.mul_coordinateSum_mul {R : Type u} [Semiring R] (shape : Shape) (left right : R) (values : Coord shape → R) : left * Semantics.coordinateSum shape values * right = Semantics.coordinateSum shape fun (coordinate : Coord shape) => left * values coordinate * right

Move contraction-invariant factors from both ends of a coordinate sum.

Multiplication order is unchanged, so this law applies to noncommutative semirings as well as ordinary numeric scalar types.

def TorchLean.Tensor.Internal.Lowering.einsumInputProduct {R : Type u} [Storage R] [Mul R] [OfNat R 1] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) : Coord checked.output → Coord (List.map checked.axisLength checked.contractedAxes) → R

The generic scalar product used by einsumTensor.

The public elaborator may replace this function by code generated from the checked literal. Its correctness obligation is pointwise equality with this reference implementation, so every generated loop reuses the same semantic bridge.

theorem TorchLean.Tensor.Internal.Lowering.einsumInputProduct_eq_foldl {R : Type u} [Storage R] [Mul R] [OfNat R 1] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) (outputCoordinate : Coord checked.output) (contractionCoordinate : Coord (List.map checked.axisLength checked.contractedAxes)) : einsumInputProduct checked inputTensors outputCoordinate contractionCoordinate = Fin.foldl checked.inputShapes.length (fun (product : R) (operand : Fin checked.inputShapes.length) => product * (inputTensors operand).getFlat (einsumInputFlatIndex checked operand outputCoordinate contractionCoordinate)) 1

The generic scalar product is the ordered finite fold of the certified input reads.

Generated kernels use this theorem as their proof boundary: optimized index arithmetic is certified operand by operand, then the resulting scalar product is compared with this fold without unfolding einsumInputProduct in a client module.

def TorchLean.Tensor.Internal.Lowering.einsumOutput {R : Type u} [Storage R] [Add R] [Mul R] [OfNat R 0] [OfNat R 1] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) : Fin checked.output.size → R

Evaluate every contracted coordinate for one flat output position.

This reference executor is completely general. Literal syntax compiles it to the same nested loops after reducing the checked pattern, which removes coordinate-pair construction from the native scalar loop.

@[inline]

def TorchLean.Tensor.Internal.Lowering.einsumTensor {R : Type u} [Storage R] [Add R] [Mul R] [OfNat R 0] [OfNat R 1] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) : checked.OutputTensor R

Execute a checked einsum using the general verified output function.

For each output entry, the executor enumerates only contracted-axis coordinates, reads the corresponding values directly from every operand, preserves source order during multiplication, and accumulates the products.

@[inline]

def TorchLean.Tensor.Internal.Lowering.einsumTensorKernel {R : Type u} [Storage R] [Add R] [Mul R] [OfNat R 0] [OfNat R 1] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) (outputValues : Fin checked.output.size → R) (_hOutputValues : ∀ (outputIndex : Fin checked.output.size), outputValues outputIndex = einsumOutput checked inputTensors outputIndex) (outputTensor : checked.OutputTensor R) (_hOutputTensor : outputTensor = Rep.ofFlatFn outputValues) : checked.OutputTensor R

Accept a compiler-generated einsum result through an explicit proof boundary.

The elaborator supplies a specialized scalar function and an output tensor, which may have been assembled sequentially or from certified parallel chunks. Erased certificates identify both with the general executor. Keeping this boundary explicit avoids optional-argument wrappers inside generated proof terms.

theorem TorchLean.Tensor.Internal.Lowering.einsumTensorKernel_correct {R : Type u} [Storage R] [Add R] [Mul R] [OfNat R 0] [OfNat R 1] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) (outputValues : Fin checked.output.size → R) (hOutputValues : ∀ (outputIndex : Fin checked.output.size), outputValues outputIndex = einsumOutput checked inputTensors outputIndex) (outputTensor : checked.OutputTensor R) (hOutputTensor : outputTensor = Rep.ofFlatFn outputValues) : einsumTensorKernel checked inputTensors outputValues hOutputValues outputTensor hOutputTensor = Semantics.denoteEinsum checked inputTensors

The compiler-generated einsum kernel equals the independent contraction denotation for every checked pattern.

No algebraic laws are required: the theorem compares the exact source-ordered operand fold and row-major contraction fold executed by both sides.

theorem TorchLean.Tensor.Internal.Lowering.einsumTensor_correct {R : Type u} [Storage R] [Add R] [Mul R] [OfNat R 0] [OfNat R 1] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) : einsumTensor checked inputTensors = Semantics.denoteEinsum checked inputTensors

The general einsum executor equals the independent contraction denotation.