Einsum plans and executable checking

This module defines the semantic data carried by a checked einsum and checks named-axis patterns against finite tensor shapes. The executable checker implements the shape behavior of einops v0.8.2:

• an input ellipsis expands to the dimensions not named explicitly;
• ellipsis dimensions are right-aligned between operands;
• dimensions attached to the same global label broadcast when either is one;
• repeated labels inside one operand must have equal dimensions and denote a diagonal; and
• output labels must occur in an input.

The mathematical semantics accepts any finite number of named labels. The pinned Python implementation caps distinct labels at 52 because it encodes them as ASCII letters, but that ceiling belongs to the encoding rather than to einsum, so nothing here enforces it.

The solitary name _ is an ordinary einsum label. It is not the skipped-axis wildcard used by parse_shape.

CheckedEinsum is the only einsum plan structure. It stores the parsed pattern, input shapes, and one total logical-axis length function. Its invariants prove that each physical dimension is either that resolved length or a singleton, and that every resolved length is supplied by an occurrence of the axis unless all occurrences are singleton. This representation supports symbolic Nat dimensions without adding a parallel symbolic plan.

Reference

The surface restrictions follow einops v0.8.2, pinned at commit 8e911db71f2e693a0c434b041180388c685ed06f. Its 52-label limit comes from einops.einops._compactify_pattern_for_einsum, which encodes labels as ASCII letters; it is not an einsum law. Dimension broadcasting is made explicit here because the Python front end delegates it to tensor backends.

inductive TorchLean.Tensor.Internal.Check.EinsumAxis : Type

The identity of one logical einsum dimension after ellipsis expansion.

Named labels preserve the user's full identifier. Ellipsis slots are numbered from left to right in the common right-aligned ellipsis shape.

• named (name : String) : EinsumAxis

  A logical dimension identified by the label written in the pattern.

• ellipsis (index : ℕ) : EinsumAxis

  One slot of the common right-aligned ellipsis dimensions.

@[instance_reducible]

instance TorchLean.Tensor.Internal.Check.instReprEinsumAxis : Repr EinsumAxis

def TorchLean.Tensor.Internal.Check.instReprEinsumAxis.repr : EinsumAxis → ℕ → Std.Format

def TorchLean.Tensor.Internal.Check.instBEqEinsumAxis.beq : EinsumAxis → EinsumAxis → Bool

@[implicit_reducible]

instance TorchLean.Tensor.Internal.Check.instBEqEinsumAxis : BEq EinsumAxis

instance TorchLean.Tensor.Internal.Check.instReflBEqEinsumAxis : ReflBEq EinsumAxis

instance TorchLean.Tensor.Internal.Check.instLawfulBEqEinsumAxis : LawfulBEq EinsumAxis

def TorchLean.Tensor.Internal.Check.instDecidableEqEinsumAxis.decEq (x✝ x✝¹ : EinsumAxis) : Decidable (x✝ = x✝¹)

@[implicit_reducible]

instance TorchLean.Tensor.Internal.Check.instDecidableEqEinsumAxis : DecidableEq EinsumAxis

def TorchLean.Tensor.Internal.Check.EinsumAxis.description : EinsumAxis → String

A readable label used in concrete shape diagnostics.

@[implicit_reducible]

def TorchLean.Tensor.Internal.Check.einsumExpressionSupported (expression : Syntax.Expression) : Bool

Whether an einsum expression uses only the v0.8.2-supported axis forms.

Empty expressions represent scalar tensors and are accepted. Each nonempty top-level position must contain exactly one named label or one bare ellipsis. Parentheses around a single named label are harmless, as in the reference parser, but genuine compositions, unit axes, and numeric axes are rejected.

@[implicit_reducible]

def TorchLean.Tensor.Internal.Check.einsumPatternSupported (pattern : Syntax.EinsumPattern) : Bool

Whether every input and the output use the supported einsum surface.

@[implicit_reducible]

def TorchLean.Tensor.Internal.Check.fixedEinsumRank (expression : Syntax.Expression) : ℕ

The number of non-ellipsis tensor dimensions written in an expression.

@[implicit_reducible]

def TorchLean.Tensor.Internal.Check.inputEllipsisRank (expression : Syntax.Expression) (inputShape : Shape) : ℕ

The number of physical dimensions captured by one input ellipsis.

Rank checking proves that the subtraction is exact. For an expression without an ellipsis the value is zero.

@[implicit_reducible]

def TorchLean.Tensor.Internal.Check.einsumEllipsisRank (pattern : Syntax.EinsumPattern) (inputShapes : List Shape) : ℕ

The common ellipsis rank of an einsum call.

Each operand contributes its own captured rank. Taking their maximum creates the target to which shorter ellipses are right-aligned.

@[implicit_reducible]

def TorchLean.Tensor.Internal.Check.expandedEllipsisAxes (commonRank operandRank : ℕ) : List EinsumAxis

Common ellipsis slots occupied by an operand with the given local rank.

theorem TorchLean.Tensor.Internal.Check.expandedEllipsisAxes_nodup (commonRank operandRank : ℕ) : (expandedEllipsisAxes commonRank operandRank).Nodup

Expanding an ellipsis never introduces duplicate logical axes.

@[implicit_reducible]

def TorchLean.Tensor.Internal.Check.expandEinsumExpression (commonRank ellipsisRank : ℕ) (expression : Syntax.Expression) : List EinsumAxis

Expand one supported expression to logical labels.

ellipsisRank is the local captured rank for an input and the common rank for the output. Unsupported composites contribute no labels; checked patterns prove that this fallback is unreachable.

@[implicit_reducible]

def TorchLean.Tensor.Internal.Check.einsumInputAxes (pattern : Syntax.EinsumPattern) (inputShapes : List Shape) : List (List EinsumAxis)

Expanded logical labels for every input operand.

@[implicit_reducible]

def TorchLean.Tensor.Internal.Check.einsumOutputAxes (pattern : Syntax.EinsumPattern) (inputShapes : List Shape) : List EinsumAxis

Expanded logical labels of the output tensor.

@[implicit_reducible]

def TorchLean.Tensor.Internal.Check.einsumGlobalAxes (pattern : Syntax.EinsumPattern) (inputShapes : List Shape) : List EinsumAxis

All logical labels in first-occurrence order.

Repeated input labels are collapsed here, but remain present in each operand's expanded list where they impose diagonal indexing.

@[implicit_reducible]

def TorchLean.Tensor.Internal.Check.einsumAxisOccurrences (pattern : Syntax.EinsumPattern) (inputShapes : List Shape) : List (EinsumAxis × ℕ)

Every physical input dimension paired with its expanded logical label.

@[implicit_reducible]

def TorchLean.Tensor.Internal.Check.einsumAxisLength (pattern : Syntax.EinsumPattern) (inputShapes : List Shape) (axis : EinsumAxis) : ℕ

The broadcast dimension of one logical label.

The first non-singleton occurrence determines the result; if every occurrence is singleton, the result is one. Compatibility checking proves that every other occurrence is either singleton or equal to this dimension. In particular, a zero dimension combined with singleton dimensions resolves to zero rather than one.

@[implicit_reducible]

def TorchLean.Tensor.Internal.Check.einsumOutput (pattern : Syntax.EinsumPattern) (inputShapes : List Shape) : Shape

The concrete output, in the axis order written by the user.

@[implicit_reducible]

def TorchLean.Tensor.Internal.Check.repeatedEinsumDimensionsAgree (axes : List EinsumAxis) (inputShape : Shape) : Bool

Whether repeated occurrences of a label inside one operand have equal physical dimensions.

theorem TorchLean.Tensor.Internal.Check.repeatedEinsumDimensionsAgree_eq_true_iff {axes : List EinsumAxis} {inputShape : Shape} : repeatedEinsumDimensionsAgree axes inputShape = true ↔ ∀ left ∈ axes.zip inputShape, ∀ right ∈ axes.zip inputShape, left.1 = right.1 → left.2 = right.2

The Boolean repeated-label check is exactly pairwise equality of dimensions for equal labels.

@[implicit_reducible]

def TorchLean.Tensor.Internal.Check.einsumAxesSubset (left right : List EinsumAxis) : Bool

Whether every axis in left occurs in right.

theorem TorchLean.Tensor.Internal.Check.einsumAxesSubset_eq_true_iff {left right : List EinsumAxis} : einsumAxesSubset left right = true ↔ ∀ ⦃axis : EinsumAxis⦄, axis ∈ left → axis ∈ right

The Boolean einsum-axis subset test exactly expresses list containment.

@[reducible, inline]

abbrev TorchLean.Tensor.Internal.Check.EinsumInputDimensions (pattern : Syntax.EinsumPattern) (inputShapes : List Shape) (axisLength : EinsumAxis → ℕ) : Prop

Concrete input dimensions are singleton or equal to their resolved logical dimension.

@[implicit_reducible]

def TorchLean.Tensor.Internal.Check.einsumAxisLengths (pattern : Syntax.EinsumPattern) (inputShapes : List Shape) (axisLength : EinsumAxis → ℕ) : Bool

Every resolved logical-axis length is justified by the input dimensions.

A length may be 1 when all occurrences are singleton. Otherwise one physical occurrence supplies the resolved length. Combined with EinsumInputDimensions, this rules out inventing a larger broadcast result for an all-singleton axis.

@[reducible, inline]

abbrev TorchLean.Tensor.Internal.Check.EinsumRepeatedDimensions (pattern : Syntax.EinsumPattern) (inputShapes : List Shape) : Prop

Every operand satisfies the exact repeated-label dimension rule.

structure TorchLean.Tensor.Internal.Check.CheckedEinsum : Type

An einsum call together with the invariants needed by tensor semantics and lowering.

No normalized axes or shapes are stored independently: all are computed from pattern and inputShapes, and these fields certify those computations.

• pattern : Syntax.EinsumPattern

  Parsed surface pattern accepted by the einsum checker.

• inputShapes : List Shape

  Physical operand shapes in source order.

• axisLength : EinsumAxis → ℕ

  Resolved length of each logical named or ellipsis axis.

• input_count : self.pattern.inputs.length = self.inputShapes.length

  The pattern contains exactly one input expression per operand shape.

• surface_supported : einsumPatternSupported self.pattern = true

  Every input and output expression belongs to the supported einsum grammar.

• repeated_dimensions : EinsumRepeatedDimensions self.pattern self.inputShapes

  Repeated labels within one operand always index equal-sized dimensions.

• input_dimensions : EinsumInputDimensions self.pattern self.inputShapes self.axisLength

  Every physical dimension is its logical broadcast length or a singleton.

• axis_lengths : einsumAxisLengths self.pattern self.inputShapes self.axisLength = true

  Every non-singleton logical length is witnessed by an input occurrence.

• output_axes_nodup : (einsumOutputAxes self.pattern self.inputShapes).Nodup

  Output labels are unique after ellipsis expansion.

• output_axes_known : einsumAxesSubset (einsumOutputAxes self.pattern self.inputShapes) (einsumGlobalAxes self.pattern self.inputShapes) = true

  Every output label is supplied by at least one input operand.

def TorchLean.Tensor.Internal.Check.CheckedEinsum.inputAxes (checked : CheckedEinsum) : List (List EinsumAxis)

Expanded logical labels for every checked input tensor.

def TorchLean.Tensor.Internal.Check.CheckedEinsum.outputAxes (checked : CheckedEinsum) : List EinsumAxis

Expanded logical labels of the checked output tensor.

def TorchLean.Tensor.Internal.Check.CheckedEinsum.globalAxes (checked : CheckedEinsum) : List EinsumAxis

Distinct logical labels over which a complete assignment ranges.

def TorchLean.Tensor.Internal.Check.CheckedEinsum.contractedAxes (checked : CheckedEinsum) : List EinsumAxis

Global logical axes summed out because they are absent from the output.

theorem TorchLean.Tensor.Internal.Check.CheckedEinsum.global_axes_nodup (checked : CheckedEinsum) : checked.globalAxes.Nodup

The global logical-axis list contains each label exactly once.

def TorchLean.Tensor.Internal.Check.CheckedEinsum.output (checked : CheckedEinsum) : Shape

The inferred shape of the checked output tensor.

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

Every checked output label belongs to the global input-label list.

def TorchLean.Tensor.Internal.Check.checkEinsum (pattern : Syntax.EinsumPattern) (inputShapes : List Shape) : Result CheckedEinsum

Check an einsum pattern against the concrete shapes of all input tensors.

Errors identify operand-count, surface, rank, repeated-label, output-label, and broadcast failures separately. A successful value carries the facts used by the independent contraction semantics.