gather_nd

paddle.gather_nd ( x, index, name=None ) [源代码]

该OP是 gather 的高维推广，并且支持多轴同时索引。 index 是一个K维度的张量，它可以认为是从 x 中取K-1维张量，每一个元素是一个切片：

gather_nd - 图1

显然， index.shape[-1] <= x.rank 并且输出张量的维度是 index.shape[:-1] + x.shape[index.shape[-1]:] 。

示例：

  1. 给定:
  2.     x = [[[ 0,  1,  2,  3],
  3.               [ 4,  5,  6,  7],
  4.               [ 8,  9, 10, 11]],
  5.              [[12, 13, 14, 15],
  6.               [16, 17, 18, 19],
  7.               [20, 21, 22, 23]]]
  8.     x.shape = (2, 3, 4)
  10. - 案例 1:
  11.     index = [[1]]
  13.     gather_nd(x, index)
  14.              = [x[1, :, :]]
  15.              = [[12, 13, 14, 15],
  16.                 [16, 17, 18, 19],
  17.                 [20, 21, 22, 23]]
  19. - 案例 2:
  21.     index = [[0,2]]
  22.     gather_nd(x, index)
  23.              = [x[0, 2, :]]
  24.              = [8, 9, 10, 11]
  26. - 案例 3:
  28.     index = [[1, 2, 3]]
  29.     gather_nd(x, index)
  30.              = [x[1, 2, 3]]
  31.              = [23]

参数：

  • x (Tensor) - 输入Tensor，数据类型可以是int32，int64，float32，float64, bool。

  • index (Tensor) - 输入的索引Tensor，其数据类型int32或者int64。它的维度 index.rank 必须大于1，并且 index.shape[-1] <= x.rank 。

  • name （str，可选）- 具体用法请参见 Name ，一般无需设置，默认值为None。

返回：shape为index.shape[:-1] + x.shape[index.shape[-1]:]的Tensor，数据类型与 x 一致。

代码示例：

  1. import paddle
  2. import numpy as np
  4. np_x = np.array([[[1, 2], [3, 4], [5, 6]],
  5.                  [[7, 8], [9, 10], [11, 12]]])
  6. np_index = [[0, 1]]
  7. x = paddle.to_tensor(np_x)
  8. index = paddle.to_tensor(np_index)
  10. output = paddle.gather_nd(x, index) #[[3, 4]]