TorchLean API

NN.Tensor.Internal.Elab.Einsum.Kernel.Index

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
      def TorchLean.Tensor.Internal.Elab.Impl.hoistedOutputIndexBase (inputDimensions logicalDimensions : List Lean.Expr) (outputCoordinates : List (Option Lean.Expr)) :

      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
        def TorchLean.Tensor.Internal.Elab.Impl.contractedAxisIndexContribution (contractedAxis : Check.EinsumAxis) (contractedCoordinate : Lean.Expr) (inputAxes : List Check.EinsumAxis) (inputDimensions logicalDimensions : List Lean.Expr) :

        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
          def TorchLean.Tensor.Internal.Elab.Impl.compileInputFlatIndexValue (inputAxes : List Check.EinsumAxis) (inputDimensions logicalDimensions : List Lean.Expr) (outputCoordinates : List (Option Lean.Expr)) (hoistedOutputBase : Option Lean.Expr) (hoistedContractedAxes : List Check.EinsumAxis) (contractionCoordinates : List (Check.EinsumAxis × Lean.Expr)) :

          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
              def TorchLean.Tensor.Internal.Elab.Impl.finIndexEquality (inputSize inputIndex certifiedInputIndex hIndexValue : Lean.Expr) :

              Lift equality of two bounded index values to equality of their Fin terms.

              Instances For
                def TorchLean.Tensor.Internal.Elab.Impl.compileInputRead (tensor inputSize inputIndexValue normalizedInputIndexValue certifiedInputIndex hInputIndexValue : Lean.Expr) (useNativeIndex : Bool) (directIndexBound? : Option Lean.Expr := none) (nativeIndexAssumptions : Array Lean.Expr := #[]) :

                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.

                Instances For