mindspore.ops.gather_nd

mindspore.ops.gather_nd(input_x, indices)

根据指定索引获取输入tensor的切片。

假设 indices 是一个K维的整型张量，遵循公式如下：

\[output[(i_0, ..., i_{K-2})] = input\_x[indices[(i_0, ..., i_{K-2})]]\]

需满足 \(indices.shape[-1] <= len(input\_x.shape)\) 。

参数：

• input_x (Tensor) - 输入tensor。

• indices (Tensor) - 指定索引。

返回：

Tensor

支持平台：

Ascend GPU CPU

样例：

>>> import mindspore
>>> import numpy as np
>>> input_x = mindspore.tensor([[-0.1, 0.3, 3.6], [0.4, 0.5, -3.2]], mindspore.float32)
>>> indices = mindspore.tensor([[0, 0], [1, 1]], mindspore.int32)
>>> output = mindspore.ops.gather_nd(input_x, indices)
>>> print(output)
[-0.1  0.5]