Functional

TorchLean functional helpers in the style of torch.* building blocks.

These are derived ops built from the small primitive TorchLean.Ops surface, so they work for:

• eager backend (runtime tape), and
• compiled backend (SSA/DAG via Compiled.GraphM), using the same model/loss definition.

The goal is to make losses readable without forcing users to call specialized ops like mse_loss directly.

PyTorch References

• torch.nn.functional: https://pytorch.org/docs/stable/nn.functional.html
• torch.autograd (detach/stop-grad concepts): https://pytorch.org/docs/stable/autograd.html
• torch.utils.checkpoint: https://pytorch.org/docs/stable/checkpoint.html

AD References

For background on reverse-mode AD (the idea behind tape-based autograd), see:

• Andreas Griewank and Andrea Walther, Evaluating Derivatives, 2nd ed., 2008.
• Seppo Linnainmaa, 1970 (reverse accumulation / the classic precursor to modern backprop).

Elementwise helpers

def Runtime.Autograd.TorchLean.F.listGet? {β : Type} (xs : List β) (i : ℕ) : Option β

Safe list indexing helper used in the dynamic (String-parsed) einsum/permute code paths.

def Runtime.Autograd.TorchLean.F.square {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) : m (RefTy m α s)

Elementwise square: x ↦ x * x.

PyTorch analogue: torch.square.

Checkpointing (semantics-first identity wrapper)

def Runtime.Autograd.TorchLean.F.checkpoint {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s t : Spec.Shape} (f : RefTy m α s → m (RefTy m α t)) (x : RefTy m α s) : m (RefTy m α t)

Checkpoint wrapper for API parity with PyTorch-style memory-saving patterns.

In this codebase, checkpointing is a semantic identity wrapper (checkpoint f x = f x). Backends that implement recomputation can refine this hook without changing the mathematical meaning.

Detach / stop-grad

def Runtime.Autograd.TorchLean.F.detach {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) : m (RefTy m α s)

Stop-gradient boundary (forward identity).

def Runtime.Autograd.TorchLean.F.stopGrad {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) : m (RefTy m α s)

Alias for detach.

Broadcasting helpers

def Runtime.Autograd.TorchLean.F.addB {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s₁ s₂ t : Spec.Shape} [s₁.BroadcastTo t] [s₂.BroadcastTo t] (x : RefTy m α s₁) (y : RefTy m α s₂) : m (RefTy m α t)

Broadcasting add: compute x + y after broadcasting both inputs to the target shape t.

PyTorch analogue: torch.add (broadcasting semantics).

def Runtime.Autograd.TorchLean.F.mulB {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s₁ s₂ t : Spec.Shape} [s₁.BroadcastTo t] [s₂.BroadcastTo t] (x : RefTy m α s₁) (y : RefTy m α s₂) : m (RefTy m α t)

Broadcasting multiply: compute x * y after broadcasting both inputs to the target shape t.

PyTorch analogue: torch.mul (broadcasting semantics).

Indexing helpers

def Runtime.Autograd.TorchLean.F.embedding {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {vocab dim : ℕ} (w : RefTy m α (Spec.Shape.dim vocab (Spec.Shape.dim dim Spec.Shape.scalar))) (idx : Fin vocab) : m (RefTy m α (Spec.Shape.dim dim Spec.Shape.scalar))

Embedding lookup (gather one row of an embedding table).

Given w : vocab × dim, return w[idx] : dim.

PyTorch analogue: torch.nn.functional.embedding for a single index.

Reductions

def Runtime.Autograd.TorchLean.F.mean {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) : m (RefTy m α Spec.Shape.scalar)

Mean reduction: mean(x) = sum(x) / numel(x).

PyTorch analogue: torch.mean.

Seeded RNG helpers

def Runtime.Autograd.TorchLean.F.randUniform {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (seed : ℕ) : m (RefTy m α s)

Deterministic U[0,1) tensor generator (seeded).

def Runtime.Autograd.TorchLean.F.bernoulliMask {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (keepProb : RefTy m α Spec.Shape.scalar) (seed : ℕ) : m (RefTy m α s)

Deterministic {0,1} mask generator (seeded) with scalar keep-probability input.

def Runtime.Autograd.TorchLean.F.dropoutSeeded {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) (p : α) (seed : ℕ) (training : Bool := true) : m (RefTy m α s)

Seeded dropout implemented as x * mask / keepProb where mask ∈ {0,1} is sampled from a deterministic PRNG keyed by seed.

def Runtime.Autograd.TorchLean.F.dropoutRefSeeded {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) (p : RefTy m α Spec.Shape.scalar) (seed : ℕ) (training : Bool := true) : m (RefTy m α s)

Seeded dropout where the probability is supplied as a scalar tensor ref.

This is useful in model builders where the layer definition stores p as data, avoiding an ad-hoc Float → α cast in polymorphic model code.

Einsum-ish building blocks

def Runtime.Autograd.TorchLean.F.matmul2d {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {mDim nDim pDim : ℕ} (a : RefTy m α (Spec.Shape.dim mDim (Spec.Shape.dim nDim Spec.Shape.scalar))) (b : RefTy m α (Spec.Shape.dim nDim (Spec.Shape.dim pDim Spec.Shape.scalar))) : m (RefTy m α (Spec.Shape.dim mDim (Spec.Shape.dim pDim Spec.Shape.scalar)))

Matrix matmul: [m,n] × [n,p] → [m,p].

def Runtime.Autograd.TorchLean.F.bmm {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {batch mDim nDim pDim : ℕ} (a : RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim mDim (Spec.Shape.dim nDim Spec.Shape.scalar)))) (b : RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim nDim (Spec.Shape.dim pDim Spec.Shape.scalar)))) : m (RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim mDim (Spec.Shape.dim pDim Spec.Shape.scalar))))

Batched matmul: [batch,m,n] × [batch,n,p] → [batch,m,p].

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

def Runtime.Autograd.TorchLean.F.einsumIjJkIk {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {iDim jDim kDim : ℕ} (a : RefTy m α (Spec.Shape.dim iDim (Spec.Shape.dim jDim Spec.Shape.scalar))) (b : RefTy m α (Spec.Shape.dim jDim (Spec.Shape.dim kDim Spec.Shape.scalar))) : m (RefTy m α (Spec.Shape.dim iDim (Spec.Shape.dim kDim Spec.Shape.scalar)))

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

def Runtime.Autograd.TorchLean.F.einsumBijBjkBik {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {batch iDim jDim kDim : ℕ} (a : RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim iDim (Spec.Shape.dim jDim Spec.Shape.scalar)))) (b : RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim jDim (Spec.Shape.dim kDim Spec.Shape.scalar)))) : m (RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim iDim (Spec.Shape.dim kDim Spec.Shape.scalar))))

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

def Runtime.Autograd.TorchLean.F.einsumBhidBhjdBhij {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {batch heads iDim jDim dDim : ℕ} (q : RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim heads (Spec.Shape.dim iDim (Spec.Shape.dim dDim Spec.Shape.scalar))))) (k : RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim heads (Spec.Shape.dim jDim (Spec.Shape.dim dDim Spec.Shape.scalar))))) : m (RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim heads (Spec.Shape.dim iDim (Spec.Shape.dim jDim Spec.Shape.scalar)))))

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

def Runtime.Autograd.TorchLean.F.einsumBhijBhjdBhid {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {batch heads iDim jDim dDim : ℕ} (attn : RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim heads (Spec.Shape.dim iDim (Spec.Shape.dim jDim Spec.Shape.scalar))))) (v : RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim heads (Spec.Shape.dim jDim (Spec.Shape.dim dDim Spec.Shape.scalar))))) : m (RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim heads (Spec.Shape.dim iDim (Spec.Shape.dim dDim Spec.Shape.scalar)))))

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

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

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

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

@[implicit_reducible]

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

Local decidability instance for Shape.well_formed (used by the dynamic einsum lowering).

inductive Runtime.Autograd.TorchLean.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

@[implicit_reducible]

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

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

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

@[implicit_reducible]

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

@[implicit_reducible]

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

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

structure Runtime.Autograd.TorchLean.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

  pre.

• post : List Char

  post.

• hasEll : Bool

  has Ell.

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

@[implicit_reducible]

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

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

Remove ASCII whitespace to simplify the hand-rolled parser.

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

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

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

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

• inputs : List Subscript

  inputs.

• output? : Option Subscript

@[implicit_reducible]

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

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

def Runtime.Autograd.TorchLean.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.TorchLean.F.Einsum.hasDupLabels (xs : List Label) : Bool

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

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

def Runtime.Autograd.TorchLean.F.Einsum.findIndex? (xs : List Label) (x : Label) : Option ℕ

Find the first index of a label (like List.findIdx?, but returning an Option Nat).

def Runtime.Autograd.TorchLean.F.Einsum.findIndex?.go (x : Label) (i : ℕ) : List Label → Option ℕ

def Runtime.Autograd.TorchLean.F.Einsum.swapAt {α : Type} (xs : List α) (d : ℕ) : List α

Swap adjacent elements at position d (used to implement permutations via adjacent swaps).

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

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

This mirrors the IR-side lowering strategy: represent a general permutation as a list of swaps at depths, then implement swaps with swapAdjacentAtDepth.

@[irreducible]

def Runtime.Autograd.TorchLean.F.Einsum.swapDepthsForPerm?.bubbleLeft (cur swapsRev : List ℕ) (i j : ℕ) : List ℕ × List ℕ

def Runtime.Autograd.TorchLean.F.Einsum.swapDepthsForPerm?.go (i : ℕ) (targets cur swapsRev : List ℕ) : Option (List ℕ)

def Runtime.Autograd.TorchLean.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 Shape.rank s.

def Runtime.Autograd.TorchLean.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.TorchLean.F.Einsum.Counts : Type

Occurrence counts for labels, represented as an association list.

def Runtime.Autograd.TorchLean.F.Einsum.countsFind? (cs : Counts) (x : Label) : Option ℕ

Lookup a label count.

def Runtime.Autograd.TorchLean.F.Einsum.countsFindD (cs : Counts) (x : Label) (dflt : ℕ) : ℕ

Lookup with default for Counts.

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

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

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

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

Count all labels across a list of operands.

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

Keep first occurrences of labels, preserving order.

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

@[reducible, inline]

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

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

def Runtime.Autograd.TorchLean.F.Einsum.dimFind? (mp : DimMap) (x : Label) : Option ℕ

Lookup a label’s dimension size.

def Runtime.Autograd.TorchLean.F.Einsum.dimFindD (mp : DimMap) (x : Label) (dflt : ℕ) : ℕ

Lookup with default for DimMap.

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

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

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

def Runtime.Autograd.TorchLean.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.TorchLean.F.Einsum.canBroadcastTo? (s₁ s₂ : Spec.Shape) : Option (s₁.CanBroadcastTo s₂)

Compute a Shape.CanBroadcastTo witness at runtime.

This mirrors the typeclass-based broadcasting used elsewhere, but returns an Option so the dynamic einsum lowering can fail gracefully.

def Runtime.Autograd.TorchLean.F.Einsum.permuteBySwaps {α : Type} [Context α] [DecidableEq Spec.Shape] {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.TorchLean.F.Einsum.canBroadcastRefl (s : Spec.Shape) : s.CanBroadcastTo s

Reflexive broadcast witness constructor (s can always broadcast to itself).

def Runtime.Autograd.TorchLean.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.TorchLean.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.TorchLean.F.Einsum.permForDuplicateLabels?.findUnusedIndex? (used : List ℕ) (l : Label) (i : ℕ) : List Label → Option ℕ

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

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

Build a diagonal mask tensor (spec-level) for diagonal embedding/extraction.

Given axes p and q, the resulting tensor is 1 when the indices along those axes agree, and 0 otherwise. The ip/iq parameters track the first-seen indices while recursing over dimensions.

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

Convenience wrapper around diagMaskSpecAux with fresh index-tracking state.

theorem Runtime.Autograd.TorchLean.F.Einsum.ofList_toList (s : Spec.Shape) : Spec.Shape.ofList s.toList = s

Shape.ofList is a left-inverse of Shape.toList.

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

Specialize diagMaskSpec to a concrete Shape by rewriting through Shape.toList.

def Runtime.Autograd.TorchLean.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.TorchLean.F.Einsum.firstDup?.go (seen : List (Label × ℕ)) (i : ℕ) : List Label → Option (Label × ℕ × ℕ)

theorem Runtime.Autograd.TorchLean.F.Einsum.size_appendDim_one (s : Spec.Shape) : (s.appendDim 1).size = s.size

Shape.appendDim s 1 preserves Shape.size (used to justify reshape tricks).

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

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

def Runtime.Autograd.TorchLean.F.einsumDyn {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] (equation : String) (xs : List ((s : Spec.Shape) × RefTy m α s)) : m (Option ((s : Spec.Shape) × RefTy m α s))

Runtime-checked einsum that returns an existential output shape.

Supported:

• multiple inputs, explicit/implicit output, and ellipsis (...).
• repeated labels within an operand (diagonal extraction / trace semantics).
• repeated labels in the output (diagonal embedding / zeroing off-diagonal entries).

Currently unsupported (returns none):

• non-broadcastable size mismatches.
• any case that would require gather/scatter-style indexing (not in the verifier-friendly op set).

This is implemented purely by reordering, reshaping, broadcasting, elementwise multiplication, and summing contracted axes.

def Runtime.Autograd.TorchLean.F.einsumDyn.diagonalizeOperand {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] (fuel : ℕ) (cur : (s : Spec.Shape) × RefTy m α s) (labs : List Einsum.Label) : OptionT m (((s : Spec.Shape) × RefTy m α s) × List Einsum.Label)

def Runtime.Autograd.TorchLean.F.einsumDyn.reduceContracted {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] (n : ℕ) (cur : (s : Spec.Shape) × RefTy m α s) : OptionT m ((s : Spec.Shape) × RefTy m α s)

def Runtime.Autograd.TorchLean.F.einsumDyn.go (seen : List Einsum.Label) : List Einsum.Label → List Einsum.Label

def Runtime.Autograd.TorchLean.F.einsum {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {sOut : Spec.Shape} (equation : String) (xs : List ((s : Spec.Shape) × RefTy m α s)) : m (Option (RefTy m α sOut))

einsum with an expected output shape.

Shape/axis helpers

def Runtime.Autograd.TorchLean.F.swapAdjacentAtDepth {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (depth : ℕ) (x : RefTy m α s) : m (RefTy m α (s.swapAdjacentAtDepth depth))

Swap two adjacent axes at a given nesting depth.

This is the primitive used to implement general permutations via a sequence of adjacent swaps. It corresponds to the backend op Torch.swapAdjacentAtDepth.

Core tensor semantics (PyTorch-style)

def Runtime.Autograd.TorchLean.F.hasDupNat (xs : List ℕ) : Bool

Detect duplicate Nats (used to validate axis lists at runtime).

def Runtime.Autograd.TorchLean.F.hasDupNat.go (seen : List ℕ) : List ℕ → Bool

def Runtime.Autograd.TorchLean.F.insertDesc (x : ℕ) : List ℕ → List ℕ

Insert x into a list kept in descending order.

def Runtime.Autograd.TorchLean.F.sortDesc (xs : List ℕ) : List ℕ

Sort a list of Nats in descending order (small insertion sort).

def Runtime.Autograd.TorchLean.F.moveAxisToLastSwaps (r axis : ℕ) : List ℕ

Swap depths that move an axis to the last position (for “reduce along axis” lowering).

def Runtime.Autograd.TorchLean.F.moveAxisToFrontSwaps (axis : ℕ) : List ℕ

Swap depths that move an axis to the front position.

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

Decidable Shape.well_formed for the dynamic reduction/slicing helpers.

@[implicit_reducible]

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

Local decidability instance for Shape.well_formed (used by dynamic reduction/slicing helpers).

theorem Runtime.Autograd.TorchLean.F.size_appendDim_one' (s : Spec.Shape) : (s.appendDim 1).size = s.size

Shape.appendDim s 1 preserves size (used to justify reshape in unsqueeze/keepdim code).

def Runtime.Autograd.TorchLean.F.permuteDyn {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (axes : List ℕ) (x : RefTy m α s) : m (Option ((s' : Spec.Shape) × RefTy m α s'))

Dynamic permutation: like permute, but returns an existential output shape.

PyTorch analogue: torch.permute / Tensor.permute (with runtime checks).

def Runtime.Autograd.TorchLean.F.permute {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s sOut : Spec.Shape} (axes : List ℕ) (x : RefTy m α s) : m (Option (RefTy m α sOut))

Permutation with an expected output shape.

This is a thin wrapper over permuteDyn that checks the computed shape equals sOut.

def Runtime.Autograd.TorchLean.F.Internal.reduceAlongLastSum {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] (x : (s : Spec.Shape) × RefTy m α s) : m (Option ((s' : Spec.Shape) × RefTy m α s'))

Reduce along the last axis with sum, returning the new (existential) shape.

This is the primitive step used by reduceDimsDynCore after it has permuted the requested axis to the last position.

def Runtime.Autograd.TorchLean.F.Internal.reduceAlongLastMean {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] (x : (s : Spec.Shape) × RefTy m α s) : m (Option ((s' : Spec.Shape) × RefTy m α s'))

Like reduceAlongLastSum, but using mean.

def Runtime.Autograd.TorchLean.F.Internal.reduceDimsDynCore {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] (reduceLast : (s : Spec.Shape) × RefTy m α s → m (Option ((s' : Spec.Shape) × RefTy m α s'))) {s : Spec.Shape} (axes : List ℕ) (keepdim : Bool) (x : RefTy m α s) : m (Option ((s' : Spec.Shape) × RefTy m α s'))

Core implementation for dynamic reductions over multiple axes.

This lowers “reduce along axis k” to:

1. permute axis k to the last position,
2. call reduceLast, and
3. optionally re-insert a singleton dimension when keepdim = true.

reduce_sum_dimsDyn and reduce_mean_dimsDyn are just specializations.

def Runtime.Autograd.TorchLean.F.reduceSumDimsDyn {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (axes : List ℕ) (x : RefTy m α s) (keepdim : Bool := false) : m (Option ((s' : Spec.Shape) × RefTy m α s'))

Dynamic multi-axis sum reduction (like torch.sum(x, dim=axes, keepdim=...)).

def Runtime.Autograd.TorchLean.F.reduceMeanDimsDyn {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (axes : List ℕ) (x : RefTy m α s) (keepdim : Bool := false) : m (Option ((s' : Spec.Shape) × RefTy m α s'))

Dynamic multi-axis mean reduction (like torch.mean(x, dim=axes, keepdim=...)).

def Runtime.Autograd.TorchLean.F.sliceRangeAxisDyn {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (axis start len : ℕ) (x : RefTy m α s) : m (Option ((s' : Spec.Shape) × RefTy m α s'))

Dynamic slice on an arbitrary axis.

This lowers slice_range_axisDyn axis start len to:

1. permute axis to the front,
2. call the axis-0 slice primitive, then
3. permute back.

def Runtime.Autograd.TorchLean.F.softmaxDimDyn {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (axis : ℕ) (x : RefTy m α s) : m (Option (RefTy m α s))

Dynamic softmax over an arbitrary axis (implemented by permuting to last, applying softmax, permuting back).

def Runtime.Autograd.TorchLean.F.logSoftmaxDimDyn {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (axis : ℕ) (x : RefTy m α s) (ε : α := Numbers.epsilon) : m (Option (RefTy m α s))

Dynamic log_softmax over an arbitrary axis (with optional epsilon for numerical stability).

def Runtime.Autograd.TorchLean.F.unsqueezeDyn {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (axis : ℕ) (x : RefTy m α s) : m (Option ((s' : Spec.Shape) × RefTy m α s'))

Dynamic unsqueeze: insert a singleton dimension at axis.

PyTorch analogue: torch.unsqueeze(x, dim=axis).

def Runtime.Autograd.TorchLean.F.squeezeDyn {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (axis : ℕ) (x : RefTy m α s) : m (Option ((s' : Spec.Shape) × RefTy m α s'))

Dynamic squeeze along a specific axis, requiring that axis to have size 1.

PyTorch analogue: torch.squeeze(x, dim=axis) (the dim-restricted variant).

def Runtime.Autograd.TorchLean.F.catAxis2Dyn {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] (axis : ℕ) (x y : (s : Spec.Shape) × RefTy m α s) : m (Option ((s' : Spec.Shape) × RefTy m α s'))

Dynamic concatenation of two tensors along axis (existential output shape).

This is the binary helper used by cat_axisDyn. It lowers to concat_dim0 by moving the requested axis to the front.

def Runtime.Autograd.TorchLean.F.catAxisDyn {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] (axis : ℕ) (xs : List ((s : Spec.Shape) × RefTy m α s)) : m (Option ((s' : Spec.Shape) × RefTy m α s'))

Dynamic concatenation of a list of tensors along axis (folding cat_axis2Dyn).

def Runtime.Autograd.TorchLean.F.stackAxisDyn {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] (axis : ℕ) (xs : List ((s : Spec.Shape) × RefTy m α s)) : m (Option ((s' : Spec.Shape) × RefTy m α s'))

Dynamic stack along a new axis.

PyTorch analogue: torch.stack(xs, dim=axis).

Implementation: unsqueeze each input at axis, then cat along the same axis.

def Runtime.Autograd.TorchLean.F.splitAxisDyn {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (axis : ℕ) (splitSizes : List ℕ) (x : RefTy m α s) : m (Option (List ((s' : Spec.Shape) × RefTy m α s')))

Dynamic split along an axis with explicit split sizes.

PyTorch analogue: torch.split(x, split_sizes, dim=axis).

def Runtime.Autograd.TorchLean.F.chunkAxisDyn {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (axis chunkSize : ℕ) (x : RefTy m α s) : m (Option (List ((s' : Spec.Shape) × RefTy m α s')))

Dynamic chunk along an axis, given a desired chunk size.

PyTorch analogue: torch.split(x, chunkSize, dim=axis) or torch.chunk (size-based variant).

def Runtime.Autograd.TorchLean.F.nchwToNhwc {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {n c h w : ℕ} (x : RefTy m α (Spec.Shape.dim n (Spec.Shape.dim c (Spec.Shape.dim h (Spec.Shape.dim w Spec.Shape.scalar))))) : m (RefTy m α (Spec.Shape.dim n (Spec.Shape.dim h (Spec.Shape.dim w (Spec.Shape.dim c Spec.Shape.scalar)))))

NCHW → NHWC for 4D tensors, implemented via two adjacent swaps.

def Runtime.Autograd.TorchLean.F.nhwcToNchw {α : Type} [Context α] [DecidableEq Spec.Shape] {m : Type → Type} [Monad m] [Ops m α] {n h w c : ℕ} (x : RefTy m α (Spec.Shape.dim n (Spec.Shape.dim h (Spec.Shape.dim w (Spec.Shape.dim c Spec.Shape.scalar))))) : m (RefTy m α (Spec.Shape.dim n (Spec.Shape.dim c (Spec.Shape.dim h (Spec.Shape.dim w Spec.Shape.scalar)))))

NHWC → NCHW for 4D tensors, implemented via two adjacent swaps.