oneflow.gather_nd

oneflow.gather_nd(input, index)Tensor

This operator is a high-dimensional extension of gather, index is a K-dimensional tensor, which is regarded as a index of input Tensor input.

Each element defines a slice of input:

\[output[i_{0},i_{1},...,i_{K-2}] = input[index(i_{0},i_{1},...,i_{K-2})]\]
Parameters
  • input – The input Tensor.

  • index – The slice indices.

For example:

>>> import oneflow as flow
>>> import numpy as np
>>> input = flow.tensor(np.array([[1, 2,3], [4, 5,6],[7,8,9]]), dtype=flow.float)
>>> index_1 = flow.tensor(np.array([[0], [2]]), dtype=flow.int)
>>> out_1 = flow.gather_nd(input,index_1)
>>> print(out_1.shape)
oneflow.Size([2, 3])
>>> out_1
tensor([[1., 2., 3.],
        [7., 8., 9.]], dtype=oneflow.float32)
>>> index_2 = flow.tensor(np.array([[0,2], [2,1]]), dtype=flow.int)
>>> out_2 = flow.gather_nd(input,index_2)
>>> out_2
tensor([3., 8.], dtype=oneflow.float32)