TorchLean API

NN.Tensor.Internal.Elab.Einsum.Loop

Certified loops used by generated einsum kernels #

Literal einsums execute concrete traversals with USize counters whenever the corresponding semantic length fits every Lean target. This module keeps the loop implementation and its correctness theorems separate from syntax elaboration.

Each optimized primitive has a theorem identifying it with the standard Fin.foldl, Array.ofFn, or row-major coordinate operation used by the mathematical lowering. The proof arguments erase during code generation, so the generated loop carries only its native bound, counter, and accumulator.

theorem TorchLean.Tensor.Internal.Elab.Impl.native_div_mul_add_mod (index divisor : USize) :
index / divisor * divisor + index % divisor = index

Recombining a native quotient and remainder recovers the original index.

This theorem is kept out of the global simplifier. The einsum compiler uses it only for row-major indices that it generated itself.

theorem TorchLean.Tensor.Internal.Elab.Impl.native_add_toNat_of_eq (left right : USize) (leftValue rightValue : ) (hLeft : left.toNat = leftValue) (hRight : right.toNat = rightValue) (hBound : leftValue + rightValue < USize.size) :
(left + right).toNat = leftValue + rightValue

Converting a native sum to Nat preserves a certified nonwrapping sum.

theorem TorchLean.Tensor.Internal.Elab.Impl.native_mul_toNat_of_eq (left right : USize) (leftValue rightValue : ) (hLeft : left.toNat = leftValue) (hRight : right.toNat = rightValue) (hBound : leftValue * rightValue < USize.size) :
(left * right).toNat = leftValue * rightValue

Converting a native product to Nat preserves a certified nonwrapping product.

theorem TorchLean.Tensor.Internal.Elab.Impl.native_div_toNat_of_eq (left right : USize) (leftValue rightValue : ) (hLeft : left.toNat = leftValue) (hRight : right.toNat = rightValue) :
(left / right).toNat = leftValue / rightValue

Native division respects certified equalities of both operands.

theorem TorchLean.Tensor.Internal.Elab.Impl.native_mod_toNat_of_eq (left right : USize) (leftValue rightValue : ) (hLeft : left.toNat = leftValue) (hRight : right.toNat = rightValue) :
(left % right).toNat = leftValue % rightValue

Native remainder respects certified equalities of both operands.

theorem TorchLean.Tensor.Internal.Elab.Impl.native_sub_toNat_of_eq (left right : USize) (leftValue rightValue : ) (hLeft : left.toNat = leftValue) (hRight : right.toNat = rightValue) (hLe : rightValue leftValue) :
(left - right).toNat = leftValue - rightValue

Native subtraction respects certified operand values when the subtrahend is no larger than the minuend.

theorem TorchLean.Tensor.Internal.Elab.Impl.nat_lt_of_native_lt_of_eq (left right : USize) (leftValue rightValue : ) (hLeft : left.toNat = leftValue) (hRight : right.toNat = rightValue) (hLt : left < right) :
leftValue < rightValue

A native strict comparison transports to the corresponding natural-number comparison when both operand values are certified.

theorem TorchLean.Tensor.Internal.Elab.Impl.nat_le_of_not_native_lt_of_eq (left right : USize) (leftValue rightValue : ) (hLeft : left.toNat = leftValue) (hRight : right.toNat = rightValue) (hNotLt : ¬left < right) :
rightValue leftValue

The false branch of a native strict comparison transports to the corresponding natural-number lower bound when both operand values are certified.

theorem TorchLean.Tensor.Internal.Elab.Impl.native_recombine_div_mod (index divisor product base : USize) (hProduct : product = index / divisor * divisor) :
index % divisor + (product + base) = index + base

Decode and immediately re-encode a native row-major coordinate before adding an invariant index base.

theorem TorchLean.Tensor.Internal.Elab.Impl.native_recombine_scaled_div_mod (scale stride index divisor base : USize) (hStride : stride = scale * divisor) :
scale * (index % divisor) + (stride * (index / divisor) + base) = scale * index + base

Decode and immediately re-encode a native row-major coordinate with a stride before adding an invariant index base.

This is the form produced for a flattened contraction coordinate inside an operand whose contracted axes are not the final physical dimensions.

theorem TorchLean.Tensor.Internal.Elab.Impl.native_div_lt_of_lt_mul (outer inner : ) (index innerBound : USize) (hInnerBound : innerBound.toNat = inner) (hIndex : index.toNat < outer * inner) :
(index / innerBound).toNat < outer

Dividing a bounded native product index by its row width produces a valid outer coordinate.

theorem TorchLean.Tensor.Internal.Elab.Impl.native_mod_lt_of_pos (inner : ) (index innerBound : USize) (hInnerBound : innerBound.toNat = inner) (hInner : 0 < inner) :
(index % innerBound).toNat < inner

Taking a native remainder by a positive row width produces a valid inner coordinate.

theorem TorchLean.Tensor.Internal.Elab.Impl.native_fin_div_eq_divNat (outer inner : ) (index innerBound : USize) (hInnerBound : innerBound.toNat = inner) (hIndex : index.toNat < outer * inner) :
(index / innerBound).toNat, = index.toNat, hIndex.divNat

Native quotient decoding agrees with the quotient coordinate of a finite row-major product index.

theorem TorchLean.Tensor.Internal.Elab.Impl.native_fin_mod_eq_modNat (outer inner : ) (index innerBound : USize) (hInnerBound : innerBound.toNat = inner) (hIndex : index.toNat < outer * inner) (hInner : 0 < inner) :
(index % innerBound).toNat, = index.toNat, hIndex.modNat

Native remainder decoding agrees with the remainder coordinate of a finite row-major product index.

theorem TorchLean.Tensor.Internal.Elab.Impl.fin_foldl_project {State : Type u_1} {Value : Type u_2} (length : ) (project : StateValue) (stateStep : StateFin lengthState) (valueStep : ValueFin lengthValue) (initial : State) (hStep : ∀ (state : State) (index : Fin length), valueStep (project state) index = project (stateStep state index)) :
project (Fin.foldl length stateStep initial) = Fin.foldl length valueStep (project initial)

Projecting the state of a finite fold is equivalent to folding the projected state whenever one step commutes with the projection.

def TorchLean.Tensor.Internal.Elab.Impl.coordinateFoldl {α : Type u} (shape : Shape) :
(αCoord shapeα)αα

Fold over a tensor shape in row-major coordinate order.

This definition is the proof-level reference for generated loop nests. The einsum compiler emits the same recursion directly, choosing a native counter independently for each concrete axis.

Instances For
    theorem TorchLean.Tensor.Internal.Elab.Impl.fin_foldl_four {α : Type u} (step : αFin 4α) (initial : α) :
    Fin.foldl 4 step initial = step (step (step (step initial 0) 1) 2) 3

    A four-element finite fold is the corresponding sequence of four updates.

    theorem TorchLean.Tensor.Internal.Elab.Impl.fin_foldl_eight {α : Type u} (step : αFin 8α) (initial : α) :
    Fin.foldl 8 step initial = step (step (step (step (step (step (step (step initial 0) 1) 2) 3) 4) 5) 6) 7

    An eight-element finite fold is the corresponding sequence of eight updates.

    theorem TorchLean.Tensor.Internal.Elab.Impl.toUSize_mkDivMod {blocks width : } (block : Fin blocks) (lane : Fin width) :
    USize.ofNat (block.mkDivMod lane) = USize.ofNat width * USize.ofNat block + USize.ofNat lane

    Converting a row-major product coordinate to USize exposes the native multiplication and addition used to construct it.

    The semantic certificate retains Fin.mkDivMod; generated indices use the equivalent word-sized arithmetic.

    theorem TorchLean.Tensor.Internal.Elab.Impl.fin_foldl_tiles {α : Type u} (blocks width tail : ) (step : αFin (blocks * width + tail)α) (initial : α) :
    Fin.foldl tail (fun (value : α) (index : Fin tail) => step value (Fin.natAdd (blocks * width) index)) (Fin.foldl blocks (fun (value : α) (block : Fin blocks) => Fin.foldl width (fun (value : α) (lane : Fin width) => step value (Fin.castLE (block.mkDivMod lane))) value) initial) = Fin.foldl (blocks * width + tail) step initial

    Fixed-width blocks followed by a tail enumerate exactly the original finite index interval. The theorem changes only loop nesting; it preserves the order in which step observes indices.

    theorem TorchLean.Tensor.Internal.Elab.Impl.coordinateFoldl_eq_fin_foldl {α : Type u} (shape : Shape) (step : αCoord shapeα) (initial : α) :
    coordinateFoldl shape step initial = Fin.foldl shape.size (fun (value : α) (index : Fin shape.size) => step value (Coord.unlinearize index)) initial

    The recursive coordinate fold is exactly the standard flat finite fold.

    This theorem supplies the arbitrary-rank certificate for generated output loops. In particular, the empty shape executes once and any zero-length axis executes no leaves.

    theorem TorchLean.Tensor.Internal.Elab.Impl.coordinateFoldl_eq_initial_of_size_eq_zero {α : Type u} (shape : Shape) (step : αCoord shapeα) (initial : α) (hSize : shape.size = 0) :
    coordinateFoldl shape step initial = initial

    A shape containing no coordinates leaves a fold accumulator unchanged.

    theorem TorchLean.Tensor.Internal.Elab.Impl.coordinateFoldl_project {State : Type u_1} {Value : Type u_2} (shape : Shape) (project : StateValue) (stateStep : StateCoord shapeState) (valueStep : ValueCoord shapeValue) (initial : State) (hStep : ∀ (state : State) (coordinate : Coord shape), valueStep (project state) coordinate = project (stateStep state coordinate)) :
    project (coordinateFoldl shape stateStep initial) = coordinateFoldl shape valueStep (project initial)

    Projecting the state of a row-major coordinate fold is equivalent to folding the projected state when every coordinate update commutes with the projection.

    theorem TorchLean.Tensor.Internal.Elab.Impl.coordinateFoldl_push_linearized_eq_array_ofFn {α : Type u} (shape : Shape) (values : Fin shape.sizeα) :
    coordinateFoldl shape (fun (output : Array α) (coordinate : Coord shape) => output.push (values coordinate.linearize)) (Array.emptyWithCapacity shape.size) = Array.ofFn values

    Pushing entries through the row-major coordinate fold constructs the standard flat function array.

    theorem TorchLean.Tensor.Internal.Elab.Impl.coordinateFoldl_storagePush_linearized_toArray_eq_array_ofFn {α : Type u} [storage : Storage α] (shape : Shape) (values : Fin shape.sizeα) :
    Storage.toArray (coordinateFoldl shape (fun (output : Storage.Buffer α) (coordinate : Coord shape) => Storage.push output (values coordinate.linearize)) (Storage.emptyWithCapacity shape.size)) = Array.ofFn values

    Observing a physical-buffer coordinate fold gives the standard flat function array.

    theorem TorchLean.Tensor.Internal.Elab.Impl.fin_foldl_storagePush_rankOne_toArray_eq_array_ofFn {α : Type u} [storage : Storage α] (length : ) (values : Fin (Shape.size [length])α) :
    Storage.toArray (Fin.foldl length (fun (output : Storage.Buffer α) (index : Fin length) => Storage.push output (values (Coord.linearize (index, PUnit.unit)))) (Storage.emptyWithCapacity (Shape.size [length]))) = Array.ofFn values

    A one-axis output loop observes directly as its row-major function array.

    This specialized certificate avoids normalizing the recursive coordinate fold when a tiled output compiler already emits the single Fin.foldl explicitly.

    theorem TorchLean.Tensor.Internal.Elab.Impl.storage_toArray_eq_array_ofFn_of_rankOne_eq {α : Type u} [storage : Storage α] (length : ) (values : Fin (Shape.size [length])α) (output : Storage.Buffer α) (hOutput : output = Fin.foldl length (fun (result : Storage.Buffer α) (index : Fin length) => Storage.push result (values (Coord.linearize (index, PUnit.unit)))) (Storage.emptyWithCapacity (Shape.size [length]))) :

    Transport a certified one-axis physical-buffer loop directly to its ordinary array observation.