TorchLean API

NN.Tensor.Internal.Elab.Einsum.Output

Verified einsum output generation #

This module generates arbitrary-rank output traversals and proves that their native loops implement the independent einsum semantics.

def TorchLean.Tensor.Internal.Elab.Impl.compileEinsumOutput (checked inputTensorFamily scalarType storage : Lean.Expr) (outputLengths semanticContractedLengths contractedLengths : List Lean.Expr) (contractionCoordinateMap? hContractionCoordinateMap? hContractionOriginalNodup? hContractionPermutation? : Option Lean.Expr) (inputProduct hInputProduct : Lean.Expr) (inputFactorization? : Option (Lean.Expr × Lean.Expr × Lean.Expr × Lean.Expr)) :

Compile the contraction into one arbitrary-rank output loop nest.

Every output axis follows the same recursive lowering. Concrete portable lengths use native counters; symbolic and oversized lengths use Fin.foldl. Small concrete contractions evaluate one output at a time. A concrete final axis uses four lanes from eight contraction terms when fewer than eight output positions are available, and otherwise uses eight lanes when the live read family is modest. Multi-axis tiled contractions flatten into one native loop from 32 terms for four lanes and 128 terms for eight lanes. Untiled scalar contractions retain the nested arbitrary-rank fold. Each lane keeps the scalar reduction order, and both concrete widths advance one contraction coordinate per recursive step so their callbacks inline consistently. When the contraction is one native loop, completed tiles are appended directly to the output buffer; nested contractions retain the vector result required by their outer folds. Any incomplete tile is emitted by the scalar path. The returned certificate identifies the complete loop nest with the general Array.ofFn executor.

Instances For