Algebraic semantics for einsum

A checked einsum ranges over one finite assignment of every logical axis. Each operand reads that assignment through three coordinate operations:

1. select the labels used by the operand, retaining repeated labels;
2. use the same selected coordinate at every occurrence of a repeated label, which implements diagonal indexing; and
3. project resolved dimensions to the operand's physical dimensions, mapping singleton dimensions to their unique coordinate.

The operand values are multiplied in source order and contracted coordinates are accumulated in row-major order. These ordered semantics require only Mul, Add, and scalar zero and one, so the same denotation applies to IEEE floating-point values without asserting false associativity or commutativity instances. When the scalar operations do form additive and multiplicative monoids, a theorem identifies the ordered denotation with the usual Rep.push fiber sum. Distributivity enters only in multilinearity theorems, and commutative multiplication only in the selected-operand adjoint.

Missing leading ellipsis slots need no special case: an operand simply omits those logical axes, so its value is constant while the omitted coordinates vary. This is exactly singleton broadcasting under right-aligned ellipses.

@[reducible, inline]

abbrev TorchLean.Tensor.Internal.Check.CheckedEinsum.InputTensors (checked : CheckedEinsum) (R : Type u) [Storage R] : Type u

The heterogeneous family of input tensors accepted by a checked einsum.

@[reducible, inline]

abbrev TorchLean.Tensor.Internal.Check.CheckedEinsum.OutputTensor (checked : CheckedEinsum) (R : Type u) [Storage R] : Type u

The output tensor type inferred by a checked einsum.

theorem TorchLean.Tensor.Internal.Check.CheckedEinsum.input_axes_length (checked : CheckedEinsum) : checked.inputAxes.length = checked.inputShapes.length

Expanded operand-axis lists and input shapes have the same count.

def TorchLean.Tensor.Internal.Check.CheckedEinsum.operandAxes (checked : CheckedEinsum) (operand : Fin checked.inputShapes.length) : List EinsumAxis

Expanded logical labels used by one input operand.

theorem TorchLean.Tensor.Internal.Check.CheckedEinsum.input_axis_mem_global (checked : CheckedEinsum) (operand : Fin checked.inputShapes.length) ⦃axis : EinsumAxis⦄ : axis ∈ checked.operandAxes operand → axis ∈ checked.globalAxes

Every logical label used by an operand belongs to the global assignment.

theorem TorchLean.Tensor.Internal.Check.CheckedEinsum.operand_shape_broadcastable (checked : CheckedEinsum) (operand : Fin checked.inputShapes.length) : List.Forall₂ (fun (physicalLength logicalLength : ℕ) => physicalLength = logicalLength ∨ physicalLength = 1) (checked.inputShapes.get operand) (List.map checked.axisLength (checked.operandAxes operand))

The physical shape of an operand broadcasts to its resolved logical shape.

This theorem is the proof-level bridge from the checker's per-dimension certificate to the coordinate map used by the denotation.

def TorchLean.Tensor.Internal.Check.CheckedEinsum.globalTensorCoordinateEquiv (checked : CheckedEinsum) : Coord (List.map checked.axisLength checked.globalAxes) ≃ AxisTuple checked.axisLength checked.globalAxes

Convert a global assignment coordinate to its named logical coordinates.

def TorchLean.Tensor.Internal.Check.CheckedEinsum.outputTensorCoordinateEquiv (checked : CheckedEinsum) : Coord checked.output ≃ AxisTuple checked.axisLength checked.outputAxes

Convert an output tensor coordinate to its named logical coordinates.

def TorchLean.Tensor.Internal.Check.CheckedEinsum.outputCoordinateOfGlobal (checked : CheckedEinsum) : Coord (List.map checked.axisLength checked.globalAxes) → Coord checked.output

Retain the requested output labels from a complete logical-axis assignment.

def TorchLean.Tensor.Internal.Check.CheckedEinsum.inputCoordinateOfGlobal (checked : CheckedEinsum) (operand : Fin checked.inputShapes.length) : Coord (List.map checked.axisLength checked.globalAxes) → Coord (checked.inputShapes.get operand)

Read one operand coordinate from a complete logical-axis assignment.

Repeated labels select one shared coordinate before singleton broadcasting, so this single map covers ordinary indexing, diagonals, and broadcasting.

theorem TorchLean.Tensor.Internal.Check.CheckedEinsum.contraction_axis_mem_global (checked : CheckedEinsum) ⦃axis : EinsumAxis⦄ : axis ∈ checked.contractedAxes → axis ∈ checked.globalAxes

Every contracted logical axis belongs to the complete global-axis assignment.

noncomputable def TorchLean.Tensor.Internal.Check.CheckedEinsum.contractionFiberEquiv (checked : CheckedEinsum) (outputCoordinate : Coord checked.output) : Fiber checked.outputCoordinateOfGlobal outputCoordinate ≃ AxisTuple checked.axisLength checked.contractedAxes

One einsum contraction-fiber coordinate is exactly one assignment of every global logical axis absent from the output.

The equivalence is independent of operand count, tensor rank, repeated labels, broadcasting, and axis lengths. It therefore also covers scalar outputs, contractions of no axes, and empty fibers caused by zero-length dimensions.

@[simp]

theorem TorchLean.Tensor.Internal.Check.CheckedEinsum.contractionFiberEquiv_apply (checked : CheckedEinsum) (outputCoordinate : Coord checked.output) (globalCoordinate : Fiber checked.outputCoordinateOfGlobal outputCoordinate) : (checked.contractionFiberEquiv outputCoordinate) globalCoordinate = AxisTuple.select ⋯ (checked.globalTensorCoordinateEquiv ↑globalCoordinate)

The contraction-fiber equivalence reads the contracted axes from the complete global logical-axis assignment.

theorem TorchLean.Tensor.Internal.Check.CheckedEinsum.contraction_fiber_card (checked : CheckedEinsum) (outputCoordinate : Coord checked.output) : Fintype.card (Fiber checked.outputCoordinateOfGlobal outputCoordinate) = (List.map checked.axisLength checked.contractedAxes).prod

Every output coordinate has one complete assignment for each setting of the contracted logical axes.

The fiber cardinality is therefore the product of the contracted lengths. The formula includes scalar outputs, contractions of no axes, and empty fibers caused by a zero-length contracted axis.

def TorchLean.Tensor.Internal.Check.CheckedEinsum.reconstructedGlobalCoordinate (checked : CheckedEinsum) (outputCoordinate : Coord checked.output) (contractionCoordinate : AxisTuple checked.axisLength checked.contractedAxes) : Coord (List.map checked.axisLength checked.globalAxes)

Reconstruct a complete logical coordinate from retained output coordinates and one row-major assignment of every contracted axis.

The output axes come first only in this intermediate tuple. Selection restores the checked global-axis order before converting back to a tensor coordinate.

@[simp]

theorem TorchLean.Tensor.Internal.Check.CheckedEinsum.outputCoordinateOfGlobal_reconstructedGlobalCoordinate (checked : CheckedEinsum) (outputCoordinate : Coord checked.output) (contractionCoordinate : AxisTuple checked.axisLength checked.contractedAxes) : checked.outputCoordinateOfGlobal (checked.reconstructedGlobalCoordinate outputCoordinate contractionCoordinate) = outputCoordinate

Reconstructing a global coordinate preserves its supplied output coordinate.

@[simp]

theorem TorchLean.Tensor.Internal.Check.CheckedEinsum.contractionOf_reconstructedGlobalCoordinate (checked : CheckedEinsum) (outputCoordinate : Coord checked.output) (contractionCoordinate : AxisTuple checked.axisLength checked.contractedAxes) : AxisTuple.select ⋯ (checked.globalTensorCoordinateEquiv (checked.reconstructedGlobalCoordinate outputCoordinate contractionCoordinate)) = contractionCoordinate

Selecting contracted axes from a reconstructed coordinate recovers their values.

def TorchLean.Tensor.Internal.Semantics.coordinateSum {R : Type u} [Add R] [OfNat R 0] (shape : Shape) : (Coord shape → R) → (initial : optParam R 0) → R

Add values over a multidimensional coordinate space in row-major order.

The initial accumulator is explicit because generated kernels may continue a partially computed contraction. No associativity or commutativity law is assumed; the nesting order is part of the denotation.

The fused lowering in Lowering.Einsum runs this very function rather than a private copy of it, so the executable kernel and the denotation can never drift apart in their traversal order.

theorem TorchLean.Tensor.Internal.Semantics.coordinateSum_eq_add_sum {R : Type u} [AddCommMonoid R] (shape : Shape) (values : Coord shape → R) (initial : R) : coordinateSum shape values initial = initial + ∑ coordinate : Coord shape, values coordinate

Nested row-major coordinate addition equals the finite sum for lawful addition.

This is the only place the library crosses from the ordered loop to a big operator; the lowering module reuses it instead of reproving it.

def TorchLean.Tensor.Internal.Semantics.einsumProductTensor {R : Type u} [Storage R] [Mul R] [OfNat R 1] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) : Rep R (List.map checked.axisLength checked.globalAxes)

The tensor of ordered operand products over complete logical assignments.

List.ofFn enumerates operand indices from left to right, and List.foldl records that exact multiplication order without assuming monoid laws.

theorem TorchLean.Tensor.Internal.Semantics.einsumProductTensor_get_ordered {R : Type u} [Storage R] [Mul R] [OfNat R 1] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) (globalCoordinate : Coord (List.map checked.axisLength checked.globalAxes)) : (einsumProductTensor checked inputTensors).get globalCoordinate = List.foldl (fun (x1 x2 : R) => x1 * x2) 1 (List.ofFn fun (operand : Fin checked.inputShapes.length) => (inputTensors operand).get (checked.inputCoordinateOfGlobal operand globalCoordinate))

Reading the product tensor exposes the source-ordered multiplicative fold.

@[simp]

theorem TorchLean.Tensor.Internal.Semantics.einsumProductTensor_get {R : Type u} [Storage R] [Monoid R] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) (globalCoordinate : Coord (List.map checked.axisLength checked.globalAxes)) : (einsumProductTensor checked inputTensors).get globalCoordinate = (List.ofFn fun (operand : Fin checked.inputShapes.length) => (inputTensors operand).get (checked.inputCoordinateOfGlobal operand globalCoordinate)).prod

For a lawful monoid, the source-ordered fold is the usual list product.

theorem TorchLean.Tensor.Internal.Semantics.einsumProductTensor_update_add {R : Type u} [Storage R] [Semiring R] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) (operand : Fin checked.inputShapes.length) (leftTensor rightTensor : Rep R (checked.inputShapes.get operand)) : einsumProductTensor checked (Function.update inputTensors operand (leftTensor + rightTensor)) = einsumProductTensor checked (Function.update inputTensors operand leftTensor) + einsumProductTensor checked (Function.update inputTensors operand rightTensor)

The ordered-product tensor is additive in any selected input operand.

All other operands remain fixed through Function.update. This statement covers operands of different ranks, repeated labels, singleton broadcasting, and any position in a positive-arity einsum. In particular, it does not move the selected operand through its neighbors, so a noncommutative semiring is sufficient.

def TorchLean.Tensor.Internal.Semantics.denoteEinsum {R : Type u} [Storage R] [Add R] [Mul R] [OfNat R 0] [OfNat R 1] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) : checked.OutputTensor R

Independent einsum denotation: sum ordered operand products over contraction coordinates while retaining the output labels in the user's requested order.

Both operand multiplication and contraction addition have explicit left-to-right orders. This is the semantic contract used for scalar types, such as IEEE floats, whose operations do not satisfy the algebraic laws required by unordered finite products and sums.

theorem TorchLean.Tensor.Internal.Semantics.denoteEinsum_apply_reconstructed {R : Type u} [Storage R] [AddCommMonoid R] [Monoid R] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) (outputCoordinate : Coord checked.output) (globalCoordinate : AxisTuple checked.axisLength checked.contractedAxes → Coord (List.map checked.axisLength checked.globalAxes)) (retainsOutput : ∀ (contractionCoordinate : AxisTuple checked.axisLength checked.contractedAxes), checked.outputCoordinateOfGlobal (globalCoordinate contractionCoordinate) = outputCoordinate) (recoversContraction : ∀ (contractionCoordinate : AxisTuple checked.axisLength checked.contractedAxes), AxisTuple.select ⋯ (checked.globalTensorCoordinateEquiv (globalCoordinate contractionCoordinate)) = contractionCoordinate) : Rep.get (denoteEinsum checked inputTensors) outputCoordinate = ∑ contractionCoordinate : AxisTuple checked.axisLength checked.contractedAxes, (einsumProductTensor checked inputTensors).get (globalCoordinate contractionCoordinate)

Evaluate an einsum using any explicit reconstruction of the complete global coordinate from the retained output coordinate and contracted-axis coordinates.

The premises state exactly that the reconstruction remains in the requested output fiber and recovers every supplied contracted-axis tuple. This theorem is independent of rank and operand count, making it suitable for relating symbolic einsum expressions to established finite-sum operations.

theorem TorchLean.Tensor.Internal.Semantics.denoteEinsum_eq_push {R : Type u} [Storage R] [AddCommMonoid R] [Monoid R] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) : denoteEinsum checked inputTensors = Rep.push checked.outputCoordinateOfGlobal (einsumProductTensor checked inputTensors)

For lawful additive and multiplicative monoids, the ordered denotation is the usual fiber sum of source-ordered operand products.

def TorchLean.Tensor.Internal.Semantics.einsumOperandVjp {R : Type u} [Storage R] [CommSemiring R] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) (operand : Fin checked.inputShapes.length) (outputCotangent : checked.OutputTensor R) : Rep R (checked.inputShapes.get operand)

Pull an output cotangent back to one selected einsum operand.

For each complete logical-axis assignment, the selected operand factor is replaced by one, leaving the product of all other operands. That product is multiplied by the output cotangent at the retained output coordinate, then all contributions are pushed back to the selected operand's physical coordinate. The push accounts uniformly for contraction, singleton broadcasting, omitted ellipsis slots, and repeated-label diagonal scatter.

The value stored in inputTensors operand is deliberately ignored: an einsum is multilinear, so its partial derivative with respect to one operand depends only on the other operands. Commutative multiplication is required to express the result using the standard tensor pairing with the selected tangent as its left factor.

theorem TorchLean.Tensor.Internal.Semantics.denoteEinsum_update_add {R : Type u} [Storage R] [Semiring R] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) (operand : Fin checked.inputShapes.length) (leftTensor rightTensor : Rep R (checked.inputShapes.get operand)) : denoteEinsum checked (Function.update inputTensors operand (leftTensor + rightTensor)) = denoteEinsum checked (Function.update inputTensors operand leftTensor) + denoteEinsum checked (Function.update inputTensors operand rightTensor)

Einsum is additive in every individual operand over an arbitrary semiring.

The selected operand may have any checked shape and may occur anywhere in the heterogeneous input family. The result follows by distributing its ordered product at each complete logical assignment and then distributing the finite sum over every contraction fiber.

theorem TorchLean.Tensor.Internal.Semantics.dot_denoteEinsum_update_eq_dot_einsumOperandVjp {R : Type u} [Storage R] [CommSemiring R] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) (operand : Fin checked.inputShapes.length) (inputTangent : Rep R (checked.inputShapes.get operand)) (outputCotangent : checked.OutputTensor R) : Rep.dot (denoteEinsum checked (Function.update inputTensors operand inputTangent)) outputCotangent = inputTangent.dot (einsumOperandVjp checked inputTensors operand outputCotangent)

The selected-operand einsum VJP is adjoint to replacing that operand by an arbitrary tangent tensor.

The statement quantifies over an arbitrary checked operand position and its dependent physical shape. It therefore covers any operand count, heterogeneous ranks, repeated-label diagonals, right-aligned ellipses, singleton broadcasting, scalar outputs, and zero-length axes without operation-specific cases.

theorem TorchLean.Tensor.Internal.Semantics.sum_denoteEinsum {R : Type u} [Storage R] [AddCommMonoid R] [Monoid R] (checked : Check.CheckedEinsum) (inputTensors : checked.InputTensors R) : ∑ outputCoordinate : Coord checked.output, Rep.get (denoteEinsum checked inputTensors) outputCoordinate = ∑ globalCoordinate : Coord (List.map checked.axisLength checked.globalAxes), (einsumProductTensor checked inputTensors).get globalCoordinate

Einsum partitions complete logical assignments by their output coordinate, so summing the output recovers the sum of all ordered operand products.