# Implement tf.gather\_nd in PyTorch

**URL:** <https://discuss.pytorch.org/t/implement-tf-gather-nd-in-pytorch/37502>\
**Category:** Uncategorized\
**Created:** [February 18, 2019, 6:17am UTC](https://discuss.pytorch.org/t/implement-tf-gather-nd-in-pytorch/37502 "2019-02-18T06:17:46Z")\
**Posts on this page:** 9\
**Page:** 1

<div class="post-metadata">

**Author:** ![hzxie](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hzxie/32/4502_2.png) [@hzxie](https://discuss.pytorch.org/u/hzxie)\
**Post date:** [February 18, 2019, 6:17am UTC](https://discuss.pytorch.org/t/implement-tf-gather-nd-in-pytorch/37502/1 "2019-02-18T06:17:46Z")

</div>

There have been some discussions about the topic, but none of them solved my problem.

Existing discussion threads:

- [How to use Tf.gather\_nd in pytorch](https://discuss.pytorch.org/t/how-to-use-tf-gather-nd-in-pytorch/28271)
- [How to do the tf.gather\_nd in pytorch?](https://discuss.pytorch.org/t/how-to-do-the-tf-gather-nd-in-pytorch/6445)
- [How to implement an equivalent of tf.gather in pytorch](https://discuss.pytorch.org/t/how-to-implement-an-equivalent-of-tf-gather-in-pytorch/33130)

I am not familiar with TF. And I want to translate a TF code to PyTorch.

The TF code goes as follows:

```auto
sample_idx = tf.cast(tf.round(sample_idx), 'int32')
g_val = tf.gather_nd(sample_grid, sample_idx) 

```

where `sample_idx` is of size `[3211264, 4]` and `sample_grid` is of size `[1, 1, 32, 32, 32]`.

Moreover, there are negative values in the `sample_idx`. And I tried to print it as follows:

```auto
tensor([[0., -18., 16., 21.],
        [0., -18., 16., 20.],
        [0., -18., 16., 20.],
        ...,
        [0., 56., -0., -5.],
        [0., 56., -0., -5.],
        [0., 56., -0., -5.]], device='cuda:0')

```

And the `g_val` is of size `[3211264]` in TF.

I wonder how to do this operation with PyTorch?

Thanks for your help.

---

<div class="post-metadata">

**Author:** ![hzxie](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hzxie/32/4502_2.png) [@hzxie](https://discuss.pytorch.org/u/hzxie)\
**Post date:** [February 18, 2019, 7:05am UTC](https://discuss.pytorch.org/t/implement-tf-gather-nd-in-pytorch/37502/2 "2019-02-18T07:05:40Z")

</div>

According to TF’s API:

> 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.

Therefore, I simply add a `clamp` on the index values. And now it is working.

```auto
def gather_nd(params, indices):
    '''
    the input indices must be a 2d tensor in the form of [[a,b,..,c],...], 
    which represents the location of the elements.
    '''
    max_value = functools.reduce(operator.mul, list(params.size())) - 1
    indices = indices.t().long()
    ndim = indices.size(0)
    idx = torch.zeros_like(indices[0]).long()
    m = 1
    
    for i in range(ndim)[::-1]:
        idx += indices[i] * m 
        m *= params.size(i)

    idx[idx < 0] = 0
    idx[idx > max_value] = 0
    return torch.take(params, idx)

```

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [February 18, 2019, 11:47pm UTC](https://discuss.pytorch.org/t/implement-tf-gather-nd-in-pytorch/37502/3 "2019-02-18T23:47:13Z")

</div>

Thanks for debugging this issue!

> [@hzxie](#):
>
> According to TF’s API:
> 
> > 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.

That seems like a weird strategy to me.

---

<div class="post-metadata">

**Author:** ![hzxie](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hzxie/32/4502_2.png) [@hzxie](https://discuss.pytorch.org/u/hzxie)\
**Post date:** [February 19, 2019, 1:46am UTC](https://discuss.pytorch.org/t/implement-tf-gather-nd-in-pytorch/37502/4 "2019-02-19T01:46:38Z")

</div>

I found that the results are not the same as TF when batch size is larger than 1.

For example, the size of `sample_idx` and `sample_grid` are `[57802752, 4]` and `[6, 32, 32, 32]`, respectively.

```auto
sample_grid = torch.from_numpy(np.load('sample_grid.npy', sample_grid)).cuda()
sample_idx = torch.from_numpy(np.load('sample_idx.npy', sample_idx)).cuda()

g_val = gather_nd(sample_grid, sample_idx)
print(torch.sum(g_val[0, 0, 0, ...])) # tensor(52197., device='cuda:0')

```

However, in TF:

```auto
_sample_grid = np.load('sample_grid.npy', sample_grid)
_sample_idx = np.load('sample_idx.npy', sample_idx)

g_val = tf.gather_nd(sample_grid, sample_idx)
_g_val = sess.run(g_val, {
    sample_grid: _sample_grid,
    sample_idx: _sample_idx
})
print(np.sum(_g_val[0, 0, 0, ...])) # 43380.0

```

You can download the `sample_grid.npy` and `sample_idx.npy` from [here](https://file.io/NNKgZF).

Who can help me out ☹

---

<div class="post-metadata">

**Author:** ![hzxie](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hzxie/32/4502_2.png) [@hzxie](https://discuss.pytorch.org/u/hzxie)\
**Post date:** [February 19, 2019, 2:28am UTC](https://discuss.pytorch.org/t/implement-tf-gather-nd-in-pytorch/37502/5 "2019-02-19T02:28:23Z")

</div>

I wrote a small program to test the behavior of `TF` and the self-implemented `gather_nd` in PyTorch.

```auto
import tensorflow as tf
import torch
import numpy as np

import functools
import operator

def gather_nd(params, indices):
    '''
    the input indices must be a 2d tensor in the form of [[a,b,..,c],...], 
    which represents the location of the elements.
    '''
    max_value = functools.reduce(operator.mul, list(params.size())) - 1
    indices = indices.t().long()
    ndim = indices.size(0)
    idx = torch.zeros_like(indices[0]).long()
    m = 1
    
    for i in range(ndim)[::-1]:
        idx += indices[i] * m 
        m *= params.size(i)

    idx[idx < 0] = 0
    idx[idx > max_value] = 0
    return torch.take(params, idx)

idx = np.array([
	[0, 0, 0, 0],
	[0, 0, 1, 0],
	[0, 1, 1, 0],
	[2, 1, 1, 0],
	[1, 1, 4, 0],
])
mtx = np.reshape(range(120), (2, 3, 4, 5))
print(mtx)

tf_result = tf.gather_nd(mtx, idx)
with tf.Session() as sess:
	print('TF', sess.run(tf_result))

mtx = torch.from_numpy(mtx)
idx = torch.from_numpy(idx)
torch_result = gather_nd(mtx, idx)
print('Torch', torch_result)

```

And the output is shown as following:

```auto
TF [0 5 25 0 0]
Torch tensor([0, 5, 25, 0, 100])

```

It seem that there are something wrong with the self-implemented one.

---

<div class="post-metadata">

**Author:** ![hzxie](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hzxie/32/4502_2.png) [@hzxie](https://discuss.pytorch.org/u/hzxie)\
**Post date:** [February 19, 2019, 2:51am UTC](https://discuss.pytorch.org/t/implement-tf-gather-nd-in-pytorch/37502/6 "2019-02-19T02:51:22Z")

</div>

I preprocessed the indices before `torch.take` and make sure that the index value in each dimension is correct.

The modified `gather_nd` is listed as follows:

```auto
def gather_nd(params, indices):
    '''
    the input indices must be a 2d tensor in the form of [[a,b,..,c],...], 
    which represents the location of the elements.
    '''
    # Normalize indices values
    params_size = list(params.size())
    assert len(params_size) == indices.size(1)

    indices[indices < 0] = 0
    for idx, ps in enumerate(params_size):
        indices[indices[:, idx] >= ps] = 0

    # Generate indices
    indices = indices.t().long()
    ndim = indices.size(0)
    idx = torch.zeros_like(indices[0]).long()
    m = 1
    
    for i in range(ndim)[::-1]:
        idx += indices[i] * m 
        m *= params.size(i)

    return torch.take(params, idx)

```

---

<div class="post-metadata">

**Author:** ![qingping](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/qingping/32/18875_2.png) [@qingping](https://discuss.pytorch.org/u/qingping)\
**Post date:** [December 22, 2019, 5:50pm UTC](https://discuss.pytorch.org/t/implement-tf-gather-nd-in-pytorch/37502/7 "2019-12-22T17:50:11Z")

</div>

Hi, I modified @hzxie ‘s code as below:  
It support the case that `indices.size(1)` is not equal to `len(param_size)`

```auto
def gather_nd(params, indices):
    '''
    the input indices must be a 2d tensor in the form of [[a,b,..,c],...],
    which represents the location of the elements.
    '''
    # Normalize indices values
    params_size = list(params.size())

    assert len(indices.size()) == 2
    assert len(params_size) >= indices.size(1)

    # Generate indices
    indices = indices.t().long()
    ndim = indices.size(0)
    idx = torch.zeros_like(indices[0]).long()
    m = 1

    for i in range(ndim)[::-1]:
        idx += indices[i] * m
        m *= params.size(i)

    params = params.reshape((-1, *tuple(torch.tensor(params.size()[ndim:]))))
    return params[idx]

```

Hope that code helps someone 😃

---

<div class="post-metadata">

**Author:** ![Ren\_yf](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ren_yf/32/19132_2.png) [@Ren\_yf](https://discuss.pytorch.org/u/Ren_yf)\
**Post date:** [December 27, 2019, 6:27am UTC](https://discuss.pytorch.org/t/implement-tf-gather-nd-in-pytorch/37502/8 "2019-12-27T06:27:47Z")

</div>

I am also quite confused with torch implementation of `tf.gather_nd` at first, but I found `tf.gather_nd` seems like a variant of slicing. I have my implementation below and hope it will help someone 😄

```python
def gather_nd(params, indices):
  """params is of "n" dimensions and has size [x1, x2, x3, ..., xn], indices is of 2 dimensions and has size [num_samples, m] (m <= n)"""
  assert type(indices) == torch.Tensor
  return params[indices.transpose(0,1).long().numpy().tolist()]

```

## testing

[Tensorflow samples](https://www.tensorflow.org/versions/r1.15/api_docs/python/tf/gather_nd)

```python
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']

```

for torch

```python
>>> params = torch.rand(2,2,2)
>>> params
tensor([[[0.4685, 0.2514],
         [0.0624, 0.0797]],

        [[0.4989, 0.1414],
         [0.6970, 0.6825]]])
>>> indices = torch.Tensor([[1]])
>>> gather_nd(params, indices)
tensor([[[0.4989, 0.1414],
         [0.6970, 0.6825]]])
>>> indices = torch.Tensor([[0, 1], [1, 0]])
>>> gather_nd(params, indices)
tensor([[0.0624, 0.0797],
        [0.4989, 0.1414]])
>>> indices = torch.Tensor([[0, 0, 1], [1, 0, 1]])
>>> gather_nd(params, indices)
tensor([0.2514, 0.1414])

```

---

<div class="post-metadata">

**Author:** ![Joris\_Mollinga](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/joris_mollinga/32/30063_2.png) [@Joris\_Mollinga](https://discuss.pytorch.org/u/Joris_Mollinga)\
**Post date:** [February 4, 2021, 12:57pm UTC](https://discuss.pytorch.org/t/implement-tf-gather-nd-in-pytorch/37502/9 "2021-02-04T12:57:01Z")

</div>

I created a version that works on 4D indices, 3D params and returns a 4D output.

Example:

```auto
params is a float32 Tensor of shape [BS, seq_len, emb_size]
indices is a int64 Tensor of shape [BS, seq_len, 10, 2]
output is a float32 Tensor of shape [BS, seq_len, 10, emb_size]

```

This implementation is **very** slow and ugly. But maybe someone can think of a clever trick to speed up the `for` loop?

```auto
def torch_gather_nd(params: torch.Tensor,
                    indices: torch.Tensor) -> torch.Tensor:
    """
    Perform the tf.gather_nd on torch.Tensor. Although working, this implementation is
    quite slow and 'ugly'. You should not care to much about performance when using
    this function. I encourage you to think about how to optimize this.

    This function has not been tested properly. It has only been tested empirically
    and shown to yield identical results compared to tf.gather_nd. Still, use at your
    own risk.

    Does not support the `batch_dims` argument that tf.gather_nd does support. This is
    something for future work.

    :param params: (Tensor) - the source Tensor
    :param indices: (LongTensor) - the indices of elements to gather
    :return output: (Tensor) – the destination tensor
    """
    assert indices.dtype == torch.int64, f"indices must be torch.LongTensor, got {indices.dtype}"
    assert indices.shape[-1] <= len(params.shape), f'The last dimension of indices can be at most the rank ' \
                                                   f'of params ({len(params.shape)})'

    # Define the output shape. According to the documentation of tf.gather_nd, this is:
    # "indices.shape[:-1] + params.shape[indices.shape[-1]:]"
    output_shape = indices.shape[:-1] + params.shape[indices.shape[-1]:]

    # Initialize the output Tensor as an empty one.
    output = torch.zeros(size=output_shape, device=params.device, dtype=params.dtype)

    # indices_to_fill is a list of tuple containing the indices to fill in `output`
    indices_to_fill = list(itertools.product(*[range(x) for x in output_shape[:-1]]))

    # Loop over the indices_to_fill and fill the `output` Tensor
    for idx in indices_to_fill:
        index_value = indices[idx]

        if len(index_value.shape) == 0:
            index_value = torch.Tensor([0, index_value.item()])

        value = params[index_value.view(-1, 1).tolist()].view(-1)
        output[idx] = value

    return output

```