gather¶

nvtripy.gather(input: Tensor, dim: int, index: Tensor) → Tensor[source]¶

Gather values from the input tensor along the specified axis based on the specified indices. This behaves similarly to numpy.take().

Parameters:

input (Tensor) – [dtype=T1] The input tensor.
dim (int) – The dimension along which to gather.
index (Tensor) – [dtype=T2] A tensor of indices to gather. Values along the provided dimension of the input are effectively replaced by the values at the specified indices. If the index tensor is multi-dimensional, the values along the dimension will be replaced by tensors instead of scalars.

Returns:

[dtype=T1] A new tensor of the same shape along every dimension except dim, which will have a size equal to index.shape[0].

Return type:

Tensor

DATA TYPE CONSTRAINTS:

T1: bool, float16, bfloat16, float32, int4, int8, int32, int64
T2: int32, int64

Example

input = tp.iota((3,)) + 4
index = tp.Tensor([0, 2])
output = tp.gather(input, dim=0, index=index)

Local Variables¶

>>> input
tensor([4, 5, 6], dtype=float32, loc=gpu:0, shape=(3,))

>>> index
tensor([0, 2], dtype=int32, loc=cpu:0, shape=(2,))

>>> output
tensor([4, 6], dtype=float32, loc=gpu:0, shape=(2,))

Example: Multi-dimensional Input

input = tp.iota((3, 3)) + 4
index = tp.Tensor([0, 2])
output = tp.gather(input, dim=0, index=index)

Local Variables¶

>>> input
tensor(
    [[4, 4, 4],
     [5, 5, 5],
     [6, 6, 6]], 
    dtype=float32, loc=gpu:0, shape=(3, 3))

>>> index
tensor([0, 2], dtype=int32, loc=cpu:0, shape=(2,))

>>> output
tensor(
    [[4, 4, 4],
     [6, 6, 6]], 
    dtype=float32, loc=gpu:0, shape=(2, 3))

Example: Multi-dimensional Indices

input = tp.iota((3,)) + 4
index = tp.Tensor([[0], [2], [1], [2]])
output = tp.gather(input, dim=0, index=index)

Local Variables¶

>>> input
tensor([4, 5, 6], dtype=float32, loc=gpu:0, shape=(3,))

>>> index
tensor(
    [[0],
     [2],
     [1],
     [2]], 
    dtype=int32, loc=cpu:0, shape=(4, 1))

>>> output
tensor(
    [[4],
     [6],
     [5],
     [6]], 
    dtype=float32, loc=gpu:0, shape=(4, 1))

Example: Multi-dimensional Input And Indices

input = tp.iota((3, 3)) + 4
index = tp.Tensor([[0, 2], [2, 1]])
output = tp.gather(input, dim=0, index=index)

Local Variables¶

>>> input
tensor(
    [[4, 4, 4],
     [5, 5, 5],
     [6, 6, 6]], 
    dtype=float32, loc=gpu:0, shape=(3, 3))

>>> index
tensor(
    [[0, 2],
     [2, 1]], 
    dtype=int32, loc=cpu:0, shape=(2, 2))

>>> output
tensor(
    [[[4, 4, 4],
      [6, 6, 6]],

     [[6, 6, 6],
      [5, 5, 5]]], 
    dtype=float32, loc=gpu:0, shape=(2, 2, 3))