Einsum-ish building blocks

Typed einsum wrappers (fast, total)

These are non-Option equivalents for the most common einsum contractions in ML code. They are intended to be used directly (no string parsing), and serve as the fast-path targets for einsum?.

def Runtime.Autograd.Model.F.einsumIjJkIk {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {iDim jDim kDim : ℕ} (a : RefTy m α [iDim, jDim]) (b : RefTy m α [jDim, kDim]) : m (RefTy m α [iDim, kDim])

einsum("ij,jk->ik", A, B) as a typed matmul.

def Runtime.Autograd.Model.F.einsumBijBjkBik {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {batch iDim jDim kDim : ℕ} (a : RefTy m α [batch, iDim, jDim]) (b : RefTy m α [batch, jDim, kDim]) : m (RefTy m α [batch, iDim, kDim])

einsum("bij,bjk->bik", A, B) as a typed batched matmul.

def Runtime.Autograd.Model.F.einsumBhidBhjdBhij {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {batch heads iDim jDim dDim : ℕ} (q : RefTy m α [batch, heads, iDim, dDim]) (k : RefTy m α [batch, heads, jDim, dDim]) : m (RefTy m α [batch, heads, iDim, jDim])

Einsum pattern used in attention: bhid,bhjd -> bhij (batched Q·Kᵀ per head).

def Runtime.Autograd.Model.F.einsumBhijBhjdBhid {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {batch heads iDim jDim dDim : ℕ} (attn : RefTy m α [batch, heads, iDim, jDim]) (v : RefTy m α [batch, heads, jDim, dDim]) : m (RefTy m α [batch, heads, iDim, dDim])

Einsum pattern used in attention: bhij,bhjd -> bhid (batched Attn·V per head).

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

def Runtime.Autograd.Model.F.Einsum.wellFormedDec (s : Spec.Shape) : Decidable s.wellFormed

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

@[instance_reducible]

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

Local decidability 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

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

@[instance_reducible]

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

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

@[instance_reducible]

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

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 ...).

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).

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

@[instance_reducible]

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

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.hasDupLabels (xs : List Label) : Bool

Detect whether a list of labels contains any duplicates (order-preserving scan).

def Runtime.Autograd.Model.F.Einsum.hasDupLabels.go (seen : List Label) : List Label → Bool

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.countsFind? (cs : Counts) (x : Label) : Option ℕ

Look up how many times a parsed einsum label occurs in the current signature.

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.

def Runtime.Autograd.Model.F.Einsum.orderedUnique (xs : List Label) : List Label

Keep first occurrences of labels, preserving order.

def Runtime.Autograd.Model.F.Einsum.orderedUnique.go (seen : List Label) : List Label → List Label

@[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.dimFind? (mp : DimMap) (x : Label) : Option ℕ

Lookup a label’s dimension size.

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) × RefTy m α s) (swaps : List ℕ) : m ((s : Spec.Shape) × RefTy m α 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 : RefTy m α s) (depths : List ℕ) : m (RefTy m α (s.applyAdjacentSwaps depths))

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

def Runtime.Autograd.Model.F.Einsum.removeAt {α : Type} : List α → ℕ → List α

Remove the element at index n (0-based), leaving the list unchanged if out of bounds.

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.diagMaskSpec {α : Type} [TorchLean.Storage α] [Zero α] [One α] (dims : List ℕ) (p q : ℕ) : TorchLean.Tensor α (Spec.Shape.ofList dims)

Diagonal mask specification with fresh index-tracking state.

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

Specialize diagMaskSpec to a concrete Shape.

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).