Function gatherND

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

indices is a K-dimensional integer tensor, best thought of as a (K-1)-dimensional tensor of indices into input, where each element defines a slice of input: output[\(i_0, ..., i_{K-2}\)] = input[indices[\(i_0, ..., i_{K-2}\)]]

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

The last dimension of indices can be at most the rank of input: indices.shape[-1] <= input.rank

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

Note that on CPU, if an out of bound index is found, an error is returned. On GPU, if an out of bound index is found, a 0 is stored in the corresponding output value.

const indices = tf.tensor2d([0, 1, 1, 0], [2,2], 'int32');
const input = tf.tensor2d([9, 10, 11, 12], [2, 2]);
tf.gatherND(input, indices).print() // [10, 11]
Copy

Doc