TorchLean API

Docs HomeInstallationGuideExamplesGraphsTools

NN.Runtime.Autograd.Torch.Core.Ops.Indexing

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.

Indexing operations #

source
def Runtime.Autograd.Torch.Internal.EagerSession.gatherScalar {α : Type} (s : EagerSession α) [Zero α] [DecidableEq Spec.Shape] {n : ℕ} (x : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (i : Fin n) :
IO (TensorRef α Spec.Shape.scalar)

Gather a scalar from a 1D vector with a Fin n index. PyTorch: x[i].

Instances For
    source
    def Runtime.Autograd.Torch.Internal.EagerSession.gatherRow {α : Type} (s : EagerSession α) [Zero α] [DecidableEq Spec.Shape] {rows cols : ℕ} (x : TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (i : Fin rows) :
    IO (TensorRef α (Spec.Shape.dim cols Spec.Shape.scalar))

    Gather a row from a 2D tensor with a Fin rows index. PyTorch: x[i] for 2D tensors.

    Instances For
      source
      def Runtime.Autograd.Torch.Internal.EagerSession.gatherScalarNat {α : Type} (s : EagerSession α) [Zero α] [DecidableEq Spec.Shape] {n : ℕ} (x : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (i : ℕ) :
      IO (TensorRef α Spec.Shape.scalar)

      Gather a scalar from a 1D vector with a raw Nat index (totalized by the tape op).

      Instances For
        source
        def Runtime.Autograd.Torch.Internal.EagerSession.gatherScalarRef {α : Type} (s : EagerSession α) [Zero α] [DecidableEq Spec.Shape] {n : ℕ} (x : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (i : NatRef) :
        IO (TensorRef α Spec.Shape.scalar)

        Dynamic gather scalar using an index stored in NatRef.

        Instances For
          source
          def Runtime.Autograd.Torch.Internal.EagerSession.gatherRowRef {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Zero α] [DecidableEq Spec.Shape] {rows cols : ℕ} (x : TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (i : NatRef) :
          IO (TensorRef α (Spec.Shape.dim cols Spec.Shape.scalar))

          Dynamic gather row using an index stored in NatRef (out-of-range gives a zero row).

          Instances For
            source
            def Runtime.Autograd.Torch.Internal.EagerSession.gatherVecNat {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {n k : ℕ} (x : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (idx : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) :
            IO (TensorRef α (Spec.Shape.dim k Spec.Shape.scalar))

            Gather k scalars using an explicit index tensor. PyTorch analogue: gather / advanced indexing.

            Instances For
              source
              def Runtime.Autograd.Torch.Internal.EagerSession.gatherRowsNat {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {rows cols k : ℕ} (x : TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (idx : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) :
              IO (TensorRef α (Spec.Shape.dim k (Spec.Shape.dim cols Spec.Shape.scalar)))

              Gather k rows using an explicit index tensor. PyTorch: index_select(dim=0, index=...).

              Instances For
                source
                def Runtime.Autograd.Torch.Internal.EagerSession.gatherVecRef {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {n k : ℕ} (x : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (idx : NatVecRef k) :
                IO (TensorRef α (Spec.Shape.dim k Spec.Shape.scalar))

                Gather k scalars using indices stored in the nat-environment (NatVecRef).

                Instances For
                  source
                  def Runtime.Autograd.Torch.Internal.EagerSession.gatherRowsRef {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {rows cols k : ℕ} (x : TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (idx : NatVecRef k) :
                  IO (TensorRef α (Spec.Shape.dim k (Spec.Shape.dim cols Spec.Shape.scalar)))

                  Gather k rows using indices stored in the nat-environment (NatVecRef).

                  Instances For
                    source
                    def Runtime.Autograd.Torch.Internal.EagerSession.scatterAddVec {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {n : ℕ} (x : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (v : TensorRef α Spec.Shape.scalar) (i : Fin n) :
                    IO (TensorRef α (Spec.Shape.dim n Spec.Shape.scalar))

                    Scatter-add into a vector: return a copy of x with x[i] += v.

                    Instances For
                      source
                      def Runtime.Autograd.Torch.Internal.EagerSession.scatterAddRow {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {rows cols : ℕ} (x : TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (v : TensorRef α (Spec.Shape.dim cols Spec.Shape.scalar)) (i : Fin rows) :
                      IO (TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar)))

                      Scatter-add into a matrix row: return a copy of x with x[i,:] += v.

                      Instances For