mindspore.ops.gather_nd

mindspore.ops.gather_nd(input_x, indices)[source]

Gathers slices from the input tensor by specified indices.

Suppose indices is an K-dimensional integer tensor, follow the formula below:

o u t p u t [(i_{0}, . . ., i_{K - 2})] = i n p u t_x [i n d i c e s [(i_{0}, . . ., i_{K - 2})]]

Must be satisfied $i n d i c e s . s h a p e [- 1] <= i n p u t_x . r a n k$ .

Parameters

input_x (Tensor) – The input tensor.
indices (Tensor) – The specified indices.

Returns

Tensor

Supported Platforms:: Ascend GPU CPU

Examples

>>> 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]