gather_nd

paddle.fluid.layers.gather_nd(input, index, name=None)[源代码]

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

gather_nd - 图1

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

示例:

  1. 给定:
  2. input = [[[ 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. input.shape = (2, 3, 4)
  9. - 案例 1:
  10. index = [[1]]
  11. gather_nd(input, index)
  12. = [input[1, :, :]]
  13. = [[12, 13, 14, 15],
  14. [16, 17, 18, 19],
  15. [20, 21, 22, 23]]
  16. - 案例 2:
  17. index = [[0,2]]
  18. gather_nd(input, index)
  19. = [input[0, 2, :]]
  20. = [8, 9, 10, 11]
  21. - 案例 3:
  22. index = [[1, 2, 3]]
  23. gather_nd(input, index)
  24. = [input[1, 2, 3]]
  25. = [23]

参数

  • input (Variable) - 输入张量,数据类型可以是int32,int64,float32,float64, bool。
  • index (Variable) - 输入的索引张量,数据类型为非负int32或非负int64。它的维度 index.rank 必须大于1,并且 index.shape[-1] <= input.rank
  • name (string) - 该层的名字,默认值为None,表示会自动命名。

返回

shape为index.shape[:-1] + input.shape[index.shape[-1]:]的Tensor|LoDTensor,数据类型与 input 一致。

返回类型

Variable

代码示例

  1. import paddle.fluid as fluid
  2. x = fluid.layers.data(name='x', shape=[3, 4, 5], dtype='float32')
  3. index = fluid.layers.data(name='index', shape=[2, 2], dtype='int32')
  4. output = fluid.layers.gather_nd(x, index)