TorchLean API

Docs HomeInstallationGuideExamplesGraphsTools

NN.Runtime.Autograd.Engine.LibTorch.Ops.Indexing

CUDA Tape Operations: Concatenation, Slicing, and Indexing #

Concat / slice along dim 0 #

source
@[inline]
def Runtime.Autograd.LibTorch.Tape.concat {n m : ℕ} {s : Spec.Shape} (t : Tape SomeBuffer) (aId bId : ℕ) :
Result (Tape SomeBuffer × ℕ)

Concatenate along the leading dimension, preserving the common trailing shape.

Instances For
    source
    @[inline]
    def Runtime.Autograd.LibTorch.Tape.slice {n : ℕ} {s : Spec.Shape} (t : Tape SomeBuffer) (xId start len : ℕ) (_h : start + len ≤ n) :
    Result (Tape SomeBuffer × ℕ)

    Slice along the leading dimension: x[start:start+len].

    Instances For

      Arbitrary-axis bounded indexing #

      source
      @[reducible, inline]
      abbrev Runtime.Autograd.LibTorch.Tape.Indexing.SwapStep :
      Type

      An adjacent-swap depth paired with its checked CUDA representation.

      Instances For
        source
        def Runtime.Autograd.LibTorch.Tape.Indexing.moveAxisToFrontDepths (axis : ℕ) :
        List ℕ

        Adjacent swaps that move axis to the outermost position.

        Instances For
          source
          def Runtime.Autograd.LibTorch.Tape.Indexing.checkedSwapSteps (depths : List ℕ) :
          Result (List SwapStep)

          Check each adjacent-swap depth before it crosses the CUDA ABI.

          Instances For
            source
            def Runtime.Autograd.LibTorch.Tape.Indexing.permuteOwned (x : Buffer) (s : Spec.Shape) :
            List SwapStep → Buffer

            Apply swaps to an owned buffer, releasing each consumed intermediate.

            Instances For
              source
              def Runtime.Autograd.LibTorch.Tape.Indexing.permuteBorrowed (x : Buffer) (s : Spec.Shape) :
              List SwapStep → Buffer

              Apply swaps to a borrowed buffer, owning only the buffers created by the swaps.

              Instances For
                source
                def Runtime.Autograd.LibTorch.Tape.Indexing.releasePermutedThen (steps : List SwapStep) (permuted keep : Buffer) :
                Buffer

                Release a permuted buffer exactly when permuteBorrowed allocated it.

                Instances For
                  source
                  def Runtime.Autograd.LibTorch.Tape.Indexing.finTensorToIndexArray {n count : ℕ} (indices : TorchLean.Tensor (Fin n) [count]) :
                  Array ℕ

                  Convert bounded tensor indices to the host array expected by CUDA row kernels.

                  Instances For
                    source
                    @[inline]
                    def Runtime.Autograd.LibTorch.Tape.select {s : Spec.Shape} (t : Tape SomeBuffer) (xId axis : ℕ) [Spec.Shape.AxisInBounds axis s] (index : Fin (s.axisSize axis)) :
                    Result (Tape SomeBuffer × ℕ)

                    Select one bounded coordinate from any tensor axis.

                    Instances For
                      source
                      @[inline]
                      def Runtime.Autograd.LibTorch.Tape.indexSelect {s : Spec.Shape} (t : Tape SomeBuffer) (xId axis count : ℕ) [Spec.Shape.AxisInBounds axis s] (indices : TorchLean.Tensor (Fin (s.axisSize axis)) [count]) :
                      Result (Tape SomeBuffer × ℕ)

                      Select several bounded coordinates from any tensor axis.

                      Instances For
                        source
                        @[inline]
                        def Runtime.Autograd.LibTorch.Tape.scatterAdd {s : Spec.Shape} (t : Tape SomeBuffer) (baseId sourceId axis count : ℕ) [Spec.Shape.AxisInBounds axis s] (indices : TorchLean.Tensor (Fin (s.axisSize axis)) [count]) :
                        Result (Tape SomeBuffer × ℕ)

                        Add indexed source slices into any tensor axis.

                        Instances For