Certified einsum operand indexing #
This module generates row-major operand indices, hoists reusable coordinate contributions, and certifies direct tensor reads against the checked einsum index plan.
Hoist multi-term output-coordinate stride products out of contraction loops.
Lean's native compiler does not reliably perform loop-invariant motion for boxed natural-number arithmetic. Precomputing a multi-term base therefore removes several products and additions from every contracted iteration. Passing a one-term base through the generated loop is measurably slower than recomputing that product, so singleton dimensions and strides that reduce to zero or one are ignored and at least two remaining terms are required.
Instances For
Add one physical-axis contribution to a row-major input index.
A literal singleton contributes nothing. When the physical length is
definitionally the logical axis length, the coordinate already belongs to
that same Fin type and is necessarily zero in the singleton case, so no
runtime broadcast branch is needed. Only genuinely ambiguous symbolic
broadcasts retain the equality test.
Instances For
Combine one operand's loop-invariant output contributions in row-major index order.
Unknown physical dimensions retain the singleton-broadcast guard outside the contraction loop. The generated index certificate treats natural-number addition as associative and commutative, so output contributions can be regrouped without changing the verified flat index.
Instances For
Collect one contracted logical axis's physical row-major contribution to an operand index.
Repeated labels contribute once per physical occurrence. Symbolic singleton dimensions retain the runtime broadcast guard required by the verified plan.
Instances For
Generate one operand's physical row-major index while retaining the exact addition order of the generic verified plan.
For operands selected by shouldHoistOutputIndexTerms, all loop-invariant
output contributions arrive as one precomputed base. Contracted axes already
included in a staged base are omitted from the remaining dynamic index.
Instances For
Transport an index bound across an equality of its natural-number value.
Generated input views use this helper before composing a logical operand index with a certified source-index map.
Instances For
Lift equality of two bounded index values to equality of their Fin terms.
Instances For
Compile one operand read from a generated row-major index.
Both branches end at the same certified Fin index. Concrete buffers use
native word arithmetic and Array.uget; symbolic buffers retain ordinary
natural-number arithmetic. The native branch additionally proves that
USize.toNat recovers the generated reference index exactly, so wrapping
cannot affect execution.