gather_nd¶

paddle. gather_nd ( x, index, name=None ) [source]

This function is actually a high-dimensional extension of gather and supports for simultaneous indexing by multiple axes. index is a K-dimensional integer tensor, which is regarded as a (K-1)-dimensional tensor of index into input, where each element defines a slice of params:

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

Obviously, index.shape[-1] <= input.rank . And, the output tensor has shape index.shape[:-1] + input.shape[index.shape[-1]:] .

Given:
    x =  [[[ 0,  1,  2,  3],
           [ 4,  5,  6,  7],
           [ 8,  9, 10, 11]],
          [[12, 13, 14, 15],
           [16, 17, 18, 19],
           [20, 21, 22, 23]]]
    x.shape = (2, 3, 4)

* Case 1:
    index = [[1]]

    gather_nd(x, index)
             = [x[1, :, :]]
             = [[12, 13, 14, 15],
                [16, 17, 18, 19],
                [20, 21, 22, 23]]

* Case 2:
    index = [[0,2]]

    gather_nd(x, index)
             = [x[0, 2, :]]
             = [8, 9, 10, 11]

* Case 3:
    index = [[1, 2, 3]]

    gather_nd(x, index)
             = [x[1, 2, 3]]
             = [23]

Parameters

• x (Tensor) – The input Tensor which it’s data type should be bool, float16, float32, float64, int32, int64.

• index (Tensor) – The index input with rank > 1, index.shape[-1] <= input.rank. Its dtype should be int32, int64.

• name (str, optional) – Name for the operation (optional, default is None). For more information, please refer to Name.

Returns

-1] + input.shape[index.shape[-1]:]

Return type

output (Tensor), A tensor with the shape index.shape[

Examples

>>> import paddle

>>> x = paddle.to_tensor([[[1, 2], [3, 4], [5, 6]],
...                       [[7, 8], [9, 10], [11, 12]]])
>>> index = paddle.to_tensor([[0, 1]])

>>> output = paddle.gather_nd(x, index)
>>> print(output)
Tensor(shape=[1, 2], dtype=int64, place=Place(cpu), stop_gradient=True,
[[3, 4]])