TorchLean API

NN.Runtime.Autograd.Engine.Core.Neural

Neural-network operations for the eager engine.

This file records layer normalization, batch normalization, and multi-head attention nodes using the corresponding specification-level forward and backward operations.

def Runtime.Autograd.Tape.softmaxMatmul {α : Type} [TorchLean.Storage α] [Context α] {batch : Spec.Shape} {rows keys width : ℕ} (t : Tape (Spec.SomeTensor α)) (scoresId valueId multiplierId : ℕ) (mask : Option (TorchLean.Tensor Bool [rows, keys]) := none) :

Normalized matrix product whose tape retains the original value contraction in backward.

Instances For
    @[inline]
    def Runtime.Autograd.Tape.layerNorm {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {seqLen embedDim : ℕ} (h_seq_pos : seqLen > 0) (h_embed_pos : embedDim > 0) (t : Tape (Spec.SomeTensor α)) (xId gammaId betaId : ℕ) (epsilon : α := TorchLean.normalizationEpsilon) :

    Layer normalization for (seqLen, embedDim) tensors.

    This records a single node whose backward returns gradients for x, gamma, and beta. PyTorch comparison: torch.nn.LayerNorm(embedDim) (applied per token) / functional.layer_norm.

    Instances For
      @[inline]
      def Runtime.Autograd.Tape.batchNorm {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {channels : ℕ} {sSpatial : Spec.Shape} (hWellFormed : (Spec.Shape.dim channels sSpatial).wellFormed) (t : Tape (Spec.SomeTensor α)) (xId gammaId betaId : ℕ) (epsilon : α := TorchLean.normalizationEpsilon) :

      Batch normalization over every spatial axis of a channel-first tensor.

      Instances For
        @[inline]
        def Runtime.Autograd.Tape.attention {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {n numHeads dModel headDim : ℕ} (h1 : n ≠ 0) (t : Tape (Spec.SomeTensor α)) (wqId wkId wvId woId xId : ℕ) (mask : Option (TorchLean.Tensor Bool [n, n]) := none) :

        Multi-head self-attention.

        This is a shape-specialized attention primitive used by transformer-style models. It depends on an optional boolean (n,n) mask and returns the attended output of shape (n,dModel).

        PyTorch comparison: similar to torch.nn.MultiheadAttention / scaled dot-product attention.

        Instances For