TorchLean API

Docs HomeInstallationGuideExamplesGraphsTools

NN.Runtime.Autograd.Engine.Core.Indexing

Core Tape Indexing Operations #

This file implements gather and scatter-style tape nodes. The forward rules expose typed indexing operations, and the backward rules route upstream gradients back to the selected source coordinates.

source
def Runtime.Autograd.Tape.select {α : Type} [TorchLean.Storage α] [Zero α] {s : Spec.Shape} (t : Tape α) (xId axis : ℕ) [Spec.Shape.AxisInBounds axis s] (index : Fin (s.axisSize axis)) :
Result (Tape α × ℕ)

Select one bounded coordinate from any tensor axis.

Instances For
    source
    def Runtime.Autograd.Tape.indexSelect {α : Type} [TorchLean.Storage α] [Add α] [Zero α] {s : Spec.Shape} (t : Tape α) (xId axis count : ℕ) [Spec.Shape.AxisInBounds axis s] (indices : TorchLean.Tensor (Fin (s.axisSize axis)) [count]) :
    Result (Tape α × ℕ)

    Select several bounded coordinates from any tensor axis.

    Instances For
      source
      def Runtime.Autograd.Tape.scatterAdd {α : Type} [TorchLean.Storage α] [Add α] [Zero α] {s : Spec.Shape} (t : Tape α) (baseId sourceId axis count : ℕ) [Spec.Shape.AxisInBounds axis s] (indices : TorchLean.Tensor (Fin (s.axisSize axis)) [count]) :
      Result (Tape α × ℕ)

      Add indexed source slices into any tensor axis.

      Instances For