General einsum (PyTorch-style subscripts; runtime-checked)

@[instance_reducible]

instance Runtime.Autograd.Model.F.Einsum.instDecidableWellFormed (s : Spec.Shape) : Decidable s.wellFormed

Decidable instance for Shape.wellFormed, used by the dynamic einsum lowering.

inductive Runtime.Autograd.Model.F.Einsum.Label : Type

Label used by the dynamic einsum parser.

chr c is a concrete axis label like 'i' or 'j', while ell k is a generated ellipsis label that stands for "some number of unnamed batch-like axes".

• chr (c : Char) : Label
• ell (k : ℕ) : Label

def Runtime.Autograd.Model.F.Einsum.instBEqLabel.beq : Label → Label → Bool

@[instance_reducible]

instance Runtime.Autograd.Model.F.Einsum.instBEqLabel : BEq Label

@[instance_reducible]

instance Runtime.Autograd.Model.F.Einsum.instDecidableEqLabel : DecidableEq Label

def Runtime.Autograd.Model.F.Einsum.instDecidableEqLabel.decEq (x✝ x✝¹ : Label) : Decidable (x✝ = x✝¹)

@[instance_reducible]

instance Runtime.Autograd.Model.F.Einsum.instReprLabel : Repr Label

def Runtime.Autograd.Model.F.Einsum.instReprLabel.repr : Label → ℕ → Std.Format

structure Runtime.Autograd.Model.F.Einsum.Subscript : Type

One operand’s subscript, split around an optional ellipsis.

For example, parsing "ab...cd" yields:

• pre = ['a','b']
• post = ['c','d']
• hasEll = true

• pre : List Char

  Index letters before the ellipsis.

• post : List Char

  Index letters after the ellipsis, empty when there is none.

• hasEll : Bool

  Whether an ellipsis was present. Without this flag "ab" and "ab..." would parse to the same pair of lists, and they mean different things once batch axes are matched up.

def Runtime.Autograd.Model.F.Einsum.instReprSubscript.repr : Subscript → ℕ → Std.Format

@[instance_reducible]

instance Runtime.Autograd.Model.F.Einsum.instReprSubscript : Repr Subscript

def Runtime.Autograd.Model.F.Einsum.stripSpaces (s : String) : String

Remove ASCII whitespace to simplify the hand-rolled parser.

def Runtime.Autograd.Model.F.Einsum.parseSubscript (raw : String) : Except String Subscript

Parse a single operand subscript (with at most one ...).

Non-whitespace characters outside the ellipsis are labels; this parser does not restrict them to PyTorch's ASCII letters.

structure Runtime.Autograd.Model.F.Einsum.Parsed : Type

Parsed einsum equation: input subscripts and an optional explicit output subscript.

• inputs : List Subscript

  One subscript per operand, in the order they appear left of ->.

• output? : Option Subscript

  Explicit output subscript, or none for the implicit form, where the output axes are the letters appearing exactly once, in alphabetical order (the NumPy convention).

@[instance_reducible]

instance Runtime.Autograd.Model.F.Einsum.instReprParsed : Repr Parsed

def Runtime.Autograd.Model.F.Einsum.instReprParsed.repr : Parsed → ℕ → Std.Format

def Runtime.Autograd.Model.F.Einsum.parseEquation (raw : String) : Except String Parsed

Parse an equation of the form "a,b->c" or "a,b" (implicit output).

def Runtime.Autograd.Model.F.Einsum.swapDepthsForPerm? (perm : List ℕ) (r : ℕ) : Option (List ℕ)

Convert a permutation of axes into a sequence of adjacent swaps.

This validates that perm has length r and only mentions axes below r, then delegates to the IR lowering strategy NN.IR.Graph.swapDepthsForPerm: a general permutation is represented as a list of swap depths, and swaps are implemented with swapAdjacentAtDepth.

def Runtime.Autograd.Model.F.Einsum.expandInputLabels (sub : Subscript) (s : Spec.Shape) (maxEll : ℕ) : Except String (List Label)

Expand an input operand’s labels to a full label list matching the operand’s rank.

If the subscript contains an ellipsis, this inserts fresh Label.ell labels so that the total label count matches Spec.Shape.rank s.

def Runtime.Autograd.Model.F.Einsum.expandOutputLabels (sub : Subscript) (maxEll : ℕ) : Except String (List Label)

Expand output labels, materializing the full ellipsis range [0..maxEll) when present.

Small association-list helpers

To keep this file dependency-light, we represent maps as association lists and use small helpers instead of Std.HashMap.

@[reducible, inline]

abbrev Runtime.Autograd.Model.F.Einsum.Counts : Type

Occurrence counts for labels, represented as an association list.

def Runtime.Autograd.Model.F.Einsum.labelCount (cs : Counts) (x : Label) : ℕ

Number of occurrences of a label. An absent label has count zero.

def Runtime.Autograd.Model.F.Einsum.countsInc (cs : Counts) (x : Label) : Counts

Increment a label’s count (inserting it if absent).

def Runtime.Autograd.Model.F.Einsum.countsInc.go (x : Label) : Counts → Counts

def Runtime.Autograd.Model.F.Einsum.labelCounts (xss : List (List Label)) : Counts

Count all labels across a list of operands.

@[reducible, inline]

abbrev Runtime.Autograd.Model.F.Einsum.DimMap : Type

Map each label to its concrete dimension size (association list).

def Runtime.Autograd.Model.F.Einsum.dimUpdate (mp : DimMap) (x : Label) (d : ℕ) : DimMap

Insert/update a label’s dimension size in a DimMap.

def Runtime.Autograd.Model.F.Einsum.dimUpdate.go (x : Label) (d : ℕ) : DimMap → DimMap

def Runtime.Autograd.Model.F.Einsum.labelDimMap (xss : List (List Label)) (shapes : List Spec.Shape) : Except String DimMap

Infer a consistent label-to-dimension map from operand label lists and operand shapes.

This implements standard einsum broadcast rules: if a label is seen with both d and 1, we keep d; if two non-1 sizes disagree, we error.

def Runtime.Autograd.Model.F.Einsum.permuteBySwaps {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] (x : (s : Spec.Shape) × Ref s) (swaps : List ℕ) : m ((s : Spec.Shape) × Ref s)

Apply a permutation expressed as adjacent swap depths to an existentially-shaped tensor.

This is the runtime “apply swaps” primitive used by both .permute and the dynamic einsum output reordering.

def Runtime.Autograd.Model.F.Einsum.permuteBySwapsTyped {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : Ref s) (depths : List ℕ) : m (Ref (s.applyAdjacentSwaps depths))

Apply adjacent swaps while retaining the resulting shape in the return type.

def Runtime.Autograd.Model.F.Einsum.permForDuplicateLabels? (src tgt : List Label) : Option (List ℕ)

Compute a permutation that maps src to tgt when duplicates are present.

This is used for the “diagonal embedding” case when output labels contain repeats: we temporarily expand the output with extra axes, then permute back to the requested (possibly-duplicated) order.

def Runtime.Autograd.Model.F.Einsum.permForDuplicateLabels?.findUnusedIndex? (used : List ℕ) (l : Label) (i : ℕ) : List Label → Option ℕ

def Runtime.Autograd.Model.F.Einsum.permForDuplicateLabels?.go (src : List Label) (used : List ℕ) (tgt : List Label) (acc : List ℕ) : Option (List ℕ)

def Runtime.Autograd.Model.F.Einsum.Internal.diagMaskSpec {α : Type} [TorchLean.Storage α] [Zero α] [One α] (dims : List ℕ) (p q : ℕ) (ip iq : Option ℕ) : TorchLean.Tensor α (Spec.Shape.ofList dims)

Recursive worker for the diagonal mask, carrying the indices seen so far on axes p and q.

Both axes are counted down as the recursion walks dims, so 0 means "this axis"; when the recursion bottoms out the two remembered indices decide the entry.

def Runtime.Autograd.Model.F.Einsum.diagMaskForShape {α : Type} [TorchLean.Storage α] [Zero α] [One α] (s : Spec.Shape) (p q : ℕ) : TorchLean.Tensor α s

Diagonal mask at a concrete shape, with fresh index-tracking state.

def Runtime.Autograd.Model.F.Einsum.firstDup? (xs : List Label) : Option (Label × ℕ × ℕ)

Return the first duplicate label in xs, along with its original and duplicate positions.

def Runtime.Autograd.Model.F.Einsum.firstDup?.go (seen : List (Label × ℕ)) (i : ℕ) : List Label → Option (Label × ℕ × ℕ)

def Runtime.Autograd.Model.F.Einsum.permMoveAxisToLast (r axis : ℕ) : List ℕ

Permutation list that moves axis to the last position (keeping relative order of others).