TorchLean API

NN.Runtime.Autograd.Model.Functional.EinsumDynamic

Dynamic einsum #

einsum? interprets a PyTorch-style subscript string at run time and lowers it to the verified op set: matmul, reorder, reshape, broadcast, multiply, sum. The dynamic lowering lives in its own module so users of the typed contractions and label bookkeeping can import Functional.Einsum without elaborating it.

def Runtime.Autograd.Model.F.einsum? {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] (equation : String) (xs : List ((s : Spec.Shape) × Ref s)) :
m (Option ((s : Spec.Shape) × Ref 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, a TorchLean extension).

Labels are not restricted to ASCII letters. Diagonal extraction and embedding multiply by a zero/one mask: off-diagonal NaN or infinity is not discarded by an indexing operation and can contaminate the result. The zeroing description applies to finite arithmetic, not all IEEE values.

Currently unsupported (returns none):

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

Matrix contractions use batch-broadcast matmul after label normalization. Other equations use reordering, reshaping, broadcasting, elementwise multiplication, and summing contracted axes.

Instances For
    def Runtime.Autograd.Model.F.einsum?.diagonalizeOperand {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] (fuel : ℕ) (cur : (s : Spec.Shape) × Ref s) (labs : List Einsum.Label) :
    Instances For
      def Runtime.Autograd.Model.F.einsum?.reduceContracted {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] (n : ℕ) (cur : (s : Spec.Shape) × Ref s) :
      OptionT m ((s : Spec.Shape) × Ref s)
      Instances For
        def Runtime.Autograd.Model.F.einsum {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {sOut : Spec.Shape} (equation : String) (xs : List ((s : Spec.Shape) × Ref s)) :
        m (Option (Ref sOut))

        einsum with an expected output shape.

        Instances For