indices is a Tensor of indices into params. The index vectors are
arranged along the last axis of indices.
This is similar to tf.gather, in which indices defines slices into the
first dimension of params. In tf.gather_nd, indices defines slices into
the first N dimensions of params, where N = indices.shape[-1].
Gathering scalars
In the simplest case the vectors in indices index the full rank of params:
If indices has a rank of K, it is helpful to think indices as a
(K-1)-dimensional tensor of indices into params.
Gathering slices
If the index vectors do not index the full rank of params then each location
in the result contains a slice of params. This example collects rows from a
matrix:
Here indices contains [2] index vectors, each with a length of 1.
The index vectors each refer to rows of the params matrix. Each
row has a shape of [3] so the output shape is [2, 3].
In this case, the relationship between the shapes is:
It is helpful to think of the results in this case as tensors-of-tensors.
The shape of the outer tensor is set by the leading dimensions of indices.
While the shape of the inner tensors is the shape of a single slice.
Batches
Additionally, both params and indices can have M leading batch
dimensions that exactly match. In this case batch_dims must be set to M.
For example, to collect one row from each of a batch of matrices you could
set the leading elements of the index vectors to be their location in the
batch:
The examples below are for the case when only indices have leading extra
dimensions. If both 'params' and 'indices' have leading batch dimensions, use
the 'batch_dims' parameter to run gather_nd in batch mode.
[null,null,["Last updated 2024-04-26 UTC."],[],[],null,["# tf.gather_nd\n\n\u003cbr /\u003e\n\n|-------------------------------------------------------------------------------------------------------------------------------|\n| [View source on GitHub](https://github.com/tensorflow/tensorflow/blob/v2.16.1/tensorflow/python/ops/array_ops.py#L5381-L5384) |\n\nGather slices from `params` into a Tensor with shape specified by `indices`. \n\n tf.gather_nd(\n params, indices, batch_dims=0, name=None\n )\n\n### Used in the notebooks\n\n| Used in the guide | Used in the tutorials |\n|-------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------|\n| - [Introduction to tensor slicing](https://www.tensorflow.org/guide/tensor_slicing) | - [Parametrized Quantum Circuits for Reinforcement Learning](https://www.tensorflow.org/quantum/tutorials/quantum_reinforcement_learning) |\n\n`indices` is a `Tensor` of indices into `params`. The index vectors are\narranged along the last axis of `indices`.\n\nThis is similar to [`tf.gather`](../tf/gather), in which `indices` defines slices into the\nfirst dimension of `params`. In [`tf.gather_nd`](../tf/gather_nd), `indices` defines slices into\nthe first `N` dimensions of `params`, where `N = indices.shape[-1]`.\n| **Caution:** 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.\n\nGathering scalars\n-----------------\n\nIn the simplest case the vectors in `indices` index the full rank of `params`: \n\n tf.gather_nd(\n indices=[[0, 0],\n [1, 1]],\n params = [['a', 'b'],\n ['c', 'd']]).numpy()\n array([b'a', b'd'], dtype=object)\n\nIn this case the result has 1-axis fewer than `indices`, and each index vector\nis replaced by the scalar indexed from `params`.\n\nIn this case the shape relationship is: \n\n index_depth = indices.shape[-1]\n assert index_depth == params.shape.rank\n result_shape = indices.shape[:-1]\n\nIf `indices` has a rank of `K`, it is helpful to think `indices` as a\n(K-1)-dimensional tensor of indices into `params`.\n\nGathering slices\n----------------\n\nIf the index vectors do not index the full rank of `params` then each location\nin the result contains a slice of params. This example collects rows from a\nmatrix: \n\n tf.gather_nd(\n indices = [[1],\n [0]],\n params = [['a', 'b', 'c'],\n ['d', 'e', 'f']]).numpy()\n array([[b'd', b'e', b'f'],\n [b'a', b'b', b'c']], dtype=object)\n\nHere `indices` contains `[2]` index vectors, each with a length of `1`.\nThe index vectors each refer to rows of the `params` matrix. Each\nrow has a shape of `[3]` so the output shape is `[2, 3]`.\n\nIn this case, the relationship between the shapes is: \n\n index_depth = indices.shape[-1]\n outer_shape = indices.shape[:-1]\n assert index_depth \u003c= params.shape.rank\n inner_shape = params.shape[index_depth:]\n output_shape = outer_shape + inner_shape\n\nIt is helpful to think of the results in this case as tensors-of-tensors.\nThe shape of the outer tensor is set by the leading dimensions of `indices`.\nWhile the shape of the inner tensors is the shape of a single slice.\n\nBatches\n-------\n\nAdditionally, both `params` and `indices` can have `M` leading batch\ndimensions that exactly match. In this case `batch_dims` must be set to `M`.\n\nFor example, to collect one row from each of a batch of matrices you could\nset the leading elements of the index vectors to be their location in the\nbatch: \n\n tf.gather_nd(\n indices = [[0, 1],\n [1, 0],\n [2, 4],\n [3, 2],\n [4, 1]],\n params=tf.zeros([5, 7, 3])).shape.as_list()\n [5, 3]\n\nThe `batch_dims` argument lets you omit those leading location dimensions\nfrom the index: \n\n tf.gather_nd(\n batch_dims=1,\n indices = [[1],\n [0],\n [4],\n [2],\n [1]],\n params=tf.zeros([5, 7, 3])).shape.as_list()\n [5, 3]\n\nThis is equivalent to caling a separate `gather_nd` for each location in the\nbatch dimensions. \n\n params=tf.zeros([5, 7, 3])\n indices=tf.zeros([5, 1])\n batch_dims = 1\n\n index_depth = indices.shape[-1]\n batch_shape = indices.shape[:batch_dims]\n assert params.shape[:batch_dims] == batch_shape\n outer_shape = indices.shape[batch_dims:-1]\n assert index_depth \u003c= params.shape.rank\n inner_shape = params.shape[batch_dims + index_depth:]\n output_shape = batch_shape + outer_shape + inner_shape\n output_shape.as_list()\n [5, 3]\n\n### More examples\n\nIndexing into a 3-tensor: \n\n tf.gather_nd(\n indices = [[1]],\n params = [[['a0', 'b0'], ['c0', 'd0']],\n [['a1', 'b1'], ['c1', 'd1']]]).numpy()\n array([[[b'a1', b'b1'],\n [b'c1', b'd1']]], dtype=object)\n\n tf.gather_nd(\n indices = [[0, 1], [1, 0]],\n params = [[['a0', 'b0'], ['c0', 'd0']],\n [['a1', 'b1'], ['c1', 'd1']]]).numpy()\n array([[b'c0', b'd0'],\n [b'a1', b'b1']], dtype=object)\n\n tf.gather_nd(\n indices = [[0, 0, 1], [1, 0, 1]],\n params = [[['a0', 'b0'], ['c0', 'd0']],\n [['a1', 'b1'], ['c1', 'd1']]]).numpy()\n array([b'b0', b'b1'], dtype=object)\n\nThe examples below are for the case when only indices have leading extra\ndimensions. If both 'params' and 'indices' have leading batch dimensions, use\nthe 'batch_dims' parameter to run gather_nd in batch mode.\n\nBatched indexing into a matrix: \n\n tf.gather_nd(\n indices = [[[0, 0]], [[0, 1]]],\n params = [['a', 'b'], ['c', 'd']]).numpy()\n array([[b'a'],\n [b'b']], dtype=object)\n\nBatched slice indexing into a matrix: \n\n tf.gather_nd(\n indices = [[[1]], [[0]]],\n params = [['a', 'b'], ['c', 'd']]).numpy()\n array([[[b'c', b'd']],\n [[b'a', b'b']]], dtype=object)\n\nBatched indexing into a 3-tensor: \n\n tf.gather_nd(\n indices = [[[1]], [[0]]],\n params = [[['a0', 'b0'], ['c0', 'd0']],\n [['a1', 'b1'], ['c1', 'd1']]]).numpy()\n array([[[[b'a1', b'b1'],\n [b'c1', b'd1']]],\n [[[b'a0', b'b0'],\n [b'c0', b'd0']]]], dtype=object)\n\n tf.gather_nd(\n indices = [[[0, 1], [1, 0]], [[0, 0], [1, 1]]],\n params = [[['a0', 'b0'], ['c0', 'd0']],\n [['a1', 'b1'], ['c1', 'd1']]]).numpy()\n array([[[b'c0', b'd0'],\n [b'a1', b'b1']],\n [[b'a0', b'b0'],\n [b'c1', b'd1']]], dtype=object)\n\n tf.gather_nd(\n indices = [[[0, 0, 1], [1, 0, 1]], [[0, 1, 1], [1, 1, 0]]],\n params = [[['a0', 'b0'], ['c0', 'd0']],\n [['a1', 'b1'], ['c1', 'd1']]]).numpy()\n array([[b'b0', b'b1'],\n [b'd0', b'c1']], dtype=object)\n\nExamples with batched 'params' and 'indices': \n\n tf.gather_nd(\n batch_dims = 1,\n indices = [[1],\n [0]],\n params = [[['a0', 'b0'],\n ['c0', 'd0']],\n [['a1', 'b1'],\n ['c1', 'd1']]]).numpy()\n array([[b'c0', b'd0'],\n [b'a1', b'b1']], dtype=object)\n\n tf.gather_nd(\n batch_dims = 1,\n indices = [[[1]], [[0]]],\n params = [[['a0', 'b0'], ['c0', 'd0']],\n [['a1', 'b1'], ['c1', 'd1']]]).numpy()\n array([[[b'c0', b'd0']],\n [[b'a1', b'b1']]], dtype=object)\n\n tf.gather_nd(\n batch_dims = 1,\n indices = [[[1, 0]], [[0, 1]]],\n params = [[['a0', 'b0'], ['c0', 'd0']],\n [['a1', 'b1'], ['c1', 'd1']]]).numpy()\n array([[b'c0'],\n [b'b1']], dtype=object)\n\nSee also [`tf.gather`](../tf/gather).\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Args ---- ||\n|--------------|---------------------------------------------------------------------------------|\n| `params` | A `Tensor`. The tensor from which to gather values. |\n| `indices` | A `Tensor`. Must be one of the following types: `int32`, `int64`. Index tensor. |\n| `name` | A name for the operation (optional). |\n| `batch_dims` | An integer or a scalar 'Tensor'. The number of batch dimensions. |\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Returns ------- ||\n|---|---|\n| A `Tensor`. Has the same type as `params`. ||\n\n\u003cbr /\u003e"]]
View source on GitHub