TorchLean API

Docs HomeInstallationGuideExamplesGraphsTools

NN.Runtime.Autograd.Torch.Core.Ops.Layers

Eager Tensor Operations #

PyTorch-style tensor operations backed by the eager CPU/CUDA tapes. These wrappers record runtime nodes, dispatch CUDA kernels when requested, and preserve the typed TensorRef surface.

Neural-network layers #

source
def Runtime.Autograd.Torch.Internal.EagerSession.linear {α : Type} [TorchLean.Storage α] (s : EagerSession α) [Inhabited α] [Add α] [Mul α] [Zero α] {inDim outDim : ℕ} (w : TensorRef α [outDim, inDim]) (b : TensorRef α [outDim]) (x : TensorRef α [inDim]) :
IO (TensorRef α [outDim])

Fully-connected linear layer y = w x + b. PyTorch: torch.nn.functional.linear.

Instances For
    source
    def Runtime.Autograd.Torch.Internal.EagerSession.mseLoss {α : Type} [TorchLean.Storage α] [TensorTransfer α] (s : EagerSession α) [Inhabited α] [Add α] [Sub α] [Mul α] [Div α] [Zero α] [One α] [NatCast α] {sh : Spec.Shape} (yhat target : TensorRef α sh) :
    IO (TensorRef α Spec.Shape.scalar)

    Mean-squared-error loss returning a scalar. PyTorch: torch.nn.functional.mse_loss.

    Instances For
      source
      def Runtime.Autograd.Torch.Internal.EagerSession.layerNorm {α : Type} [TorchLean.Storage α] (s : EagerSession α) [Context α] [TensorTransfer α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {seqLen embedDim : ℕ} (h_seq_pos : seqLen > 0) (h_embed_pos : embedDim > 0) (x : TensorRef α [seqLen, embedDim]) (gamma beta : TensorRef α [embedDim]) (epsilon : α := TorchLean.normalizationEpsilon) :
      IO (TensorRef α [seqLen, embedDim])

      Layer normalization over embedding dimension. PyTorch: nn.LayerNorm / functional.layer_norm.

      Instances For
        source
        def Runtime.Autograd.Torch.Internal.EagerSession.batchNorm {α : Type} [TorchLean.Storage α] (s : EagerSession α) [Context α] [TensorTransfer α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {channels : ℕ} {sSpatial : Spec.Shape} (hWellFormed : (Spec.Shape.dim channels sSpatial).wellFormed) (x : TensorRef α (Spec.Shape.dim channels sSpatial)) (gamma beta : TensorRef α [channels]) (epsilon : α := TorchLean.normalizationEpsilon) :
        IO (TensorRef α (Spec.Shape.dim channels sSpatial))

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

        Instances For
          source
          def Runtime.Autograd.Torch.Internal.EagerSession.multiHeadAttention {α : Type} [TorchLean.Storage α] (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {n numHeads dModel headDim : ℕ} (h1 : n ≠ 0) (wq wk wv : TensorRef α [dModel, numHeads * headDim]) (wo : TensorRef α [numHeads * headDim, dModel]) (x : TensorRef α [n, dModel]) (mask : Option (TorchLean.Tensor Bool [n, n]) := none) :
          IO (TensorRef α [n, dModel])

          Multi-head self-attention (typed, proof-friendly). PyTorch: nn.MultiheadAttention (conceptually).

          Instances For