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 a certified arbitrary-rank output buffer.

The compiler first tries task-parallel outer-axis chunks. Otherwise, concrete portable buffers with at least 1,024 entries use a flat native traversal; the remaining cases use nested output folds. Native counters require portable concrete bounds, while symbolic or oversized axes retain Fin.foldl.

Nested folds may tile the final axis according to einsumOutputTileWidth?. Each lane preserves its scalar contraction order, and incomplete tiles use the scalar path. A single native contraction may fuse its completed tile directly into the output append. Every selected path supplies an equality certificate against the general Array.ofFn reference.

Instances For