tf.raw_ops.GatherNd

Gather slices from params into a Tensor with shape specified by indices.

tf.raw_ops.GatherNd(
    params, indices, bad_indices_policy='', name=None
)

indices is a K-dimensional integer tensor, best thought of as a (K-1)-dimensional tensor of indices into params, where each element defines a slice of params:

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

Whereas in tf.gather indices defines slices into the axis dimension of params, in tf.gather_nd, indices defines slices into the first N dimensions of params, where N = indices.shape[-1].

The last dimension of indices can be at most the rank of params:

indices.shape[-1] <= params.rank

The last dimension of indices corresponds to elements (if indices.shape[-1] == params.rank) or slices (if indices.shape[-1] < params.rank) along dimension indices.shape[-1] of params. The output tensor has shape

indices.shape[:-1] + params.shape[indices.shape[-1]:]

If indices contains any out-of-bound indices, depending on bad_indices_policy, the op will either return an error or ignore the out-of-bound indices. bad_indices_policy can be one of the following values:

  1. "" or "DEFAULT": raises on CPU and ignore on GPU. This is because historically on CPU and GPU we handle errors in different ways, and for backward compatibility we keep the default behavior.
  2. "ERROR": raises error; GPU does not support this value.
  3. "IGNORE": ignore error and set the corresponding output to 0; supported on both CPU and GPU.

Some examples below.

Simple indexing into a matrix:

    indices = [[0, 0], [1, 1]]
    params = [['a', 'b'], ['c', 'd']]
    output = ['a', 'd']

Slice indexing into a matrix:

    indices = [[1], [0]]
    params = [['a', 'b'], ['c', 'd']]
    output = [['c', 'd'], ['a', 'b']]

Indexing into a 3-tensor:

    indices = [[1]]
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]
    output = [[['a1', 'b1'], ['c1', 'd1']]]


    indices = [[0, 1], [1, 0]]
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]
    output = [['c0', 'd0'], ['a1', 'b1']]


    indices = [[0, 0, 1], [1, 0, 1]]
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]
    output = ['b0', 'b1']

Batched indexing into a matrix:

    indices = [[[0, 0]], [[0, 1]]]
    params = [['a', 'b'], ['c', 'd']]
    output = [['a'], ['b']]

Batched slice indexing into a matrix:

    indices = [[[1]], [[0]]]
    params = [['a', 'b'], ['c', 'd']]
    output = [[['c', 'd']], [['a', 'b']]]

Batched indexing into a 3-tensor:

    indices = [[[1]], [[0]]]
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]
    output = [[[['a1', 'b1'], ['c1', 'd1']]],
              [[['a0', 'b0'], ['c0', 'd0']]]]

    indices = [[[0, 1], [1, 0]], [[0, 0], [1, 1]]]
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]
    output = [[['c0', 'd0'], ['a1', 'b1']],
              [['a0', 'b0'], ['c1', 'd1']]]


    indices = [[[0, 0, 1], [1, 0, 1]], [[0, 1, 1], [1, 1, 0]]]
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]
    output = [['b0', 'b1'], ['d0', 'c1']]

See also tf.gather and tf.batch_gather.

params A Tensor. The tensor from which to gather values.
indices A Tensor. Must be one of the following types: int16, int32, int64. Index tensor.
bad_indices_policy An optional string. Defaults to "".
name A name for the operation (optional).

A Tensor. Has the same type as params.