Reflected einsum index construction

This module builds and certifies the coordinate and flat-index expressions used by the generated einsum kernels.

def TorchLean.Tensor.Internal.Elab.Impl.natMulExpr (left right : Lean.Expr) : Lean.MetaM Lean.Expr

Build a natural-number product without emitting statically neutral arithmetic.

def TorchLean.Tensor.Internal.Elab.Impl.natAddExpr (left right : Lean.Expr) : Lean.MetaM Lean.Expr

Build a natural-number sum without emitting statically neutral arithmetic.

def TorchLean.Tensor.Internal.Elab.Impl.shapeSizeExpr : List Lean.Expr → Lean.MetaM Lean.Expr

Build the right-associated product used by Shape.size.

def TorchLean.Tensor.Internal.Elab.Impl.runtimeShapeSizeExpr : List Lean.Expr → Lean.MetaM Lean.Expr

Compute a runtime stride while folding literal and neutral factors.

def TorchLean.Tensor.Internal.Elab.Impl.supportsNativeInputIndex (dimensions : List Lean.Expr) : Lean.MetaM Bool

Use native index arithmetic only when an operand's complete concrete buffer fits the platform word size. Symbolic shapes retain the general Nat kernel.

def TorchLean.Tensor.Internal.Elab.Impl.hasConcreteZeroDimension (lengths : List Lean.Expr) : Lean.MetaM Bool

Report whether a reflected shape contains a dimension definitionally equal to zero. The compiler uses this only for static dead-loop elimination; symbolic dimensions continue through the fully general lowering.

def TorchLean.Tensor.Internal.Elab.Impl.nativeLoopBound? (length : Lean.Expr) : Lean.MetaM (Option (Lean.Expr × Lean.Expr))

Produce a portable native bound for a definitionally concrete loop length.

Every Lean target represents at least 2^32 native indices. Restricting this optimization to that common range keeps generated code portable across 32-bit and 64-bit targets. The returned equality is erased, while the literal USize survives as the executable loop bound. Symbolic and larger lengths retain Fin.foldl.

def TorchLean.Tensor.Internal.Elab.Impl.natOperationOperands? (direct heterogeneous : Lean.Name) (value : Lean.Expr) : Option (Lean.Expr × Lean.Expr)

Return the operands of a direct or notation-elaborated binary Nat operation.

Metaprogram-generated arithmetic can retain either kernel representation after proof-producing simplification. Recognizing both keeps later native lowering independent of that representation choice.

partial def TorchLean.Tensor.Internal.Elab.Impl.nativeIndexValueWithIntermediates (value : Lean.Expr) : Lean.MetaM (Lean.Expr × List Lean.Expr)

Translate a concrete row-major Nat expression to platform-native arithmetic and report every addition or multiplication that could wrap.

The reported natural-number expressions let callers certify intermediate smallness from their loop bounds before proving that the complete native expression has the intended value.

partial def TorchLean.Tensor.Internal.Elab.Impl.nativeIndexValue (value : Lean.Expr) : Lean.MetaM Lean.Expr

Translate a concrete row-major Nat expression to platform-native arithmetic.

Literal factors remain native constants. Dynamic leaves are converted when they reach the generated expression, while additions and multiplications run as USize operations. Division, remainder, subtraction, and unsupported forms remain one compact Nat.toUSize leaf. A separate erased certificate proves that this modular expression has exactly the original natural-number value.

This compact translation is used by ordinary einsum operands. Native operation lowerings that provide explicit loop invariants use certifiedNativeIndexValue? instead, which can certify the larger arithmetic vocabulary compositionally.

def TorchLean.Tensor.Internal.Elab.Impl.coordinateAxisExpr : ℕ → Lean.Expr → Lean.MetaM Lean.Expr

Project one axis value from the nested product representation of Coord.

Literal einsums use these direct projections inside the generated nested loops, avoiding division and remainder operations in the scalar kernel.

def TorchLean.Tensor.Internal.Elab.Impl.coordinateComponentsFromFlatIndex : List Lean.Expr → Lean.Expr → Lean.MetaM (List Lean.Expr)

Expand row-major unlinearization into bounded quotient/remainder components.

Binding these components before constructing the temporary proof-level coordinate lets native code decode each output axis once without allocating the nested Prod representation of Coord.

def TorchLean.Tensor.Internal.Elab.Impl.foldOperandIndexExpr (operandCount operandIndex : ℕ) : Lean.MetaM Lean.Expr

Construct the operand index produced by repeatedly unfolding Fin.foldl.

The first operand is 0; each later operand is one more Fin.succ around zero in the remaining finite type. Matching this canonical form lets generated scalar products be certified directly by the standard fold laws.

def TorchLean.Tensor.Internal.Elab.Impl.coordinateFromComponents : List Lean.Expr → Lean.MetaM Lean.Expr

Rebuild the nested coordinate representation from bounded components.

def TorchLean.Tensor.Internal.Elab.Impl.coordinateFromComponentsEquality : List Lean.Expr → List Lean.Expr → List Lean.Expr → Lean.MetaM Lean.Expr

Combine componentwise coordinate equalities into equality of nested Coord values.

def TorchLean.Tensor.Internal.Elab.Impl.withGeneratedLetPair (bindings : List (Lean.Name × Lean.Expr)) (body : List Lean.Expr → Lean.Elab.TermElabM (Lean.Expr × Lean.Expr)) : Lean.Elab.TermElabM (Lean.Expr × Lean.Expr)

Introduce generated values as ordinary let bindings and pass their local variables to the body in source order. Both returned expressions receive the same bindings, allowing an optimized value and its correctness proof to be constructed together.

def TorchLean.Tensor.Internal.Elab.Impl.mkLambdaPreservingLets (localExpr body : Lean.Expr) : Lean.MetaM Lean.Expr

Abstract one ordinary local while retaining the body's generated lets.

def TorchLean.Tensor.Internal.Elab.Impl.withLeadingLetPair (value : Lean.Expr) (body : Array Lean.Expr → Lean.Expr → Lean.Elab.TermElabM (Lean.Expr × Lean.Expr)) : Lean.Elab.TermElabM (Lean.Expr × Lean.Expr)

Expose the leading lets of a generated expression as local declarations.

The optimized expression and its certificate are wrapped in the same lets after compilation, preserving loop-invariant values outside nested folds. The continuation also receives the opened locals so a caller can selectively unfold generated aliases without enabling unrestricted zeta reduction.

theorem TorchLean.Tensor.Internal.Elab.Impl.fin_val_eq_zero_of_eq_one {dimension : ℕ} (coordinate : Fin dimension) (hDimension : dimension = 1) : ↑coordinate = 0

A symbolic Fin coordinate is zero when its dimension is known locally to be one. This form lets the generated-kernel certifier simplify dependent coordinates without first substituting through an entire tensor coordinate.

theorem TorchLean.Tensor.Internal.Elab.Impl.evaluateInputFlatIndexPlan_ite {condition : Prop} [Decidable condition] {outputShape contractedShape : Shape} (outputCoordinate : Coord outputShape) (contractionCoordinate : Coord contractedShape) (positivePlan negativePlan : List (Bool × ℕ × ℕ)) : Lowering.evaluateInputFlatIndexPlan outputShape contractedShape outputCoordinate contractionCoordinate (if condition then positivePlan else negativePlan) = if condition then Lowering.evaluateInputFlatIndexPlan outputShape contractedShape outputCoordinate contractionCoordinate positivePlan else Lowering.evaluateInputFlatIndexPlan outputShape contractedShape outputCoordinate contractionCoordinate negativePlan

Evaluating a conditionally selected operand plan is the same as selecting the corresponding evaluated index. This exposes only plan-selection conditions to the certifier, without distributing arbitrary applications over conditionals.

def TorchLean.Tensor.Internal.Elab.Impl.certifyCompiledKernel (description : String) (proposition : Lean.Expr) : Lean.Elab.TermElabM Lean.Expr

Prove that generated row-major arithmetic agrees with the verified generic kernel after reducing the statically known pattern and operand family.

def TorchLean.Tensor.Internal.Elab.Impl.certifyNativeIndex (proposition : Lean.Expr) : Lean.Elab.TermElabM Lean.Expr

Certify that native word arithmetic has neither wrapped nor changed a generated concrete input index.

Lean supports 32-bit and 64-bit USize; splitting that platform theorem lets the simplifier expose the corresponding modulus before omega discharges the coordinate bounds.

def TorchLean.Tensor.Internal.Elab.Impl.certifyNativeIntermediateBound (value : Lean.Expr) (assumptions : Array Lean.Expr) : Lean.Elab.TermElabM Lean.Expr

Certify that one generated addition or multiplication remains below the native word modulus under explicit loop-index hypotheses.

def TorchLean.Tensor.Internal.Elab.Impl.certifyNativeArithmeticFact (description : String) (proposition : Lean.Expr) (assumptions : Array Lean.Expr) : Lean.Elab.TermElabM Lean.Expr

Certify one generated arithmetic side condition under explicit loop-index hypotheses.

The returned proof is closed over exactly the supplied assumptions. Native subtraction uses this to establish that its segment prefix does not exceed the selected packed-axis position.

partial def TorchLean.Tensor.Internal.Elab.Impl.certifiedNativeIndexValue? (value : Lean.Expr) (assumptions : Array Lean.Expr) : Lean.Elab.TermElabM (Option (Lean.Expr × Lean.Expr))

Translate generated Nat arithmetic to native words while constructing its value theorem one operation at a time.

This compositional route is used when a caller supplies explicit loop invariants. It avoids asking one arithmetic tactic to rediscover every nested machine-word nonwrapping fact after the complete expression has been built.

def TorchLean.Tensor.Internal.Elab.Impl.certifyInputIndexBound (index inputSize : Lean.Expr) : Lean.Elab.TermElabM Lean.Expr

Prove that a concrete generated row-major index lies inside its input buffer.

The index expression itself supplies every coordinate needed by the proof. Discarding unrelated ambient locals ensures the erased certificate has the same loop dependencies as the executable index.

partial def TorchLean.Tensor.Internal.Elab.Impl.exposeFinalArgumentLets (function argument : Lean.Expr) : Lean.MetaM Lean.Expr

Move leading lets from the final argument of an application around the application itself.

The generated contraction loops have the form total + body. Exposing the leading lets of body lets the loop pass below identify which index bases are independent of the current loop coordinate.

def TorchLean.Tensor.Internal.Elab.Impl.withFoldInvariantLets (body : Lean.Expr) (foldLocals : Array Lean.Expr) (continuation : Lean.Expr → Lean.MetaM Lean.Expr) : Lean.MetaM Lean.Expr

Partition a generated leading-let chain around one contraction fold.

Lets independent of the fold move outside it even when they follow a coordinate-dependent let. A let depending on an earlier dependent local stays inside as well, so the partition accounts for transitive dependencies while preserving the original order within both groups.

Every rewrite is definitionally equal to its source. A value depending on an outer contraction coordinate therefore moves across inner folds, but remains inside the fold that introduces that coordinate.

def TorchLean.Tensor.Internal.Elab.Impl.withFoldInvariantNativeIndexExpressions (body : Lean.Expr) (foldLocals : Array Lean.Expr) (continuation : Lean.Expr → Lean.MetaM Lean.Expr) : Lean.MetaM Lean.Expr

Hoist maximal native-index fragments that are invariant under the current contraction fold.

nativeIndexValue emits only native constants and arithmetic. Lean's C compiler does not always move compound fragments of that tree out of nested loops. Binding the largest independent fragment once removes repeated stride arithmetic while preserving definitional equality. Recursing through dependent let assignments lets an outer fold subsequently hoist the nested base that is independent of its own coordinate. A bare stride product is not bound: passing a singleton base through a generated callback is slower than recomputing it in simple kernels.