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)