TorchLean API

Docs HomeInstallationGuideExamplesGraphsTools

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: reorder, reshape, broadcast, multiply, sum. It is one large monadic definition and by far the most expensive thing to elaborate in this corner of the tree, so it lives on its own rather than in Functional.Einsum. Everything that only needs the typed contractions or the label bookkeeping imports that module and does not pay for this one.

source
def Runtime.Autograd.Model.F.einsum? {α : Type} [TorchLean.Storage α] [Context α] {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.

Instances For
    source
    def Runtime.Autograd.Model.F.einsum?.diagonalizeOperand {α : Type} [TorchLean.Storage α] [Context α] {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)
    Instances For
      source
      def Runtime.Autograd.Model.F.einsum?.reduceContracted {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] (n : ℕ) (cur : (s : Spec.Shape) × RefTy m α s) :
      OptionT m ((s : Spec.Shape) × RefTy m α s)
      Instances For
        source
        def Runtime.Autograd.Model.F.einsum?.go (seen : List Einsum.Label) :
        List Einsum.Label → List Einsum.Label
        Instances For
          source
          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) × RefTy m α s)) :
          m (Option (RefTy m α sOut))

          einsum with an expected output shape.

          Instances For