TorchLean API

Docs HomeInstallationGuideExamplesGraphsTools

NN.Runtime.Autograd.Engine.LibTorch.Ops.Linear

CUDA Tape Operations: Matrix and Loss Nodes #

Linear algebra #

source
@[inline]
def Runtime.Autograd.LibTorch.Tape.Internal.matmul {m n p : ℕ} (t : Tape SomeBuffer) (aId bId : ℕ) :
Result (Tape SomeBuffer × ℕ)

Backend matrix kernel for a single pair of matrices.

Instances For
    source
    @[inline]
    def Runtime.Autograd.LibTorch.Tape.Internal.matmulFlattened {batch m n p : ℕ} (t : Tape SomeBuffer) (aId bId : ℕ) (scale : Option Float := none) :
    Result (Tape SomeBuffer × ℕ)

    Backend matrix kernel over a flattened leading shape.

    Instances For

      Linear layer / losses #

      source
      @[inline]
      def Runtime.Autograd.LibTorch.Tape.linear {outDim inDim : ℕ} (t : Tape SomeBuffer) (wId bId xId : ℕ) :
      Result (Tape SomeBuffer × ℕ)

      Linear layer: y = W·x + b with W : (outDim,inDim), x : inDim, b : outDim.

      Instances For
        source
        @[inline]
        def Runtime.Autograd.LibTorch.Tape.mseLoss {s : Spec.Shape} (t : Tape SomeBuffer) (yhatId targetId : ℕ) :
        Result (Tape SomeBuffer × ℕ)

        Mean-squared-error loss with "mean" reduction (single scalar output).

        Instances For