Implement tf.gather_nd in PyTorch

hzxie: <p>There have been some discussions about the topic, but none of them solved my problem.</p>
<p>Existing discussion threads:</p>
<ul>
<li><a href="https://discuss.pytorch.org/t/how-to-use-tf-gather-nd-in-pytorch/28271" class="inline-onebox">How to use Tf.gather_nd in pytorch</a></li>
<li><a href="https://discuss.pytorch.org/t/how-to-do-the-tf-gather-nd-in-pytorch/6445" class="inline-onebox">How to do the tf.gather_nd in pytorch?</a></li>
<li><a href="https://discuss.pytorch.org/t/how-to-implement-an-equivalent-of-tf-gather-in-pytorch/33130" class="inline-onebox">How to implement an equivalent of tf.gather in pytorch</a></li>
</ul>
<p>I am not familiar with TF. And I want to translate a TF code to PyTorch.</p>
<p>The TF code goes as follows:</p>
<pre><code class="lang-auto">sample_idx = tf.cast(tf.round(sample_idx), 'int32')
g_val = tf.gather_nd(sample_grid, sample_idx) 
</code></pre>
<p>where <code>sample_idx</code> is of size <code>[3211264, 4]</code> and <code>sample_grid</code> is of size <code>[1, 1, 32, 32, 32]</code>.</p>
<p>Moreover, there are negative values in the <code>sample_idx</code>. And I tried to print it as follows:</p>
<pre><code class="lang-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')
</code></pre>
<p>And the <code>g_val</code> is of size <code>[3211264]</code> in TF.</p>
<p>I wonder how to do this operation with PyTorch?</p>
<p>Thanks for your help.</p>

hzxie: <p>According to TF’s API:</p>
<blockquote>
<p>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.</p>
</blockquote>
<p>Therefore, I simply add a <code>clamp</code> on the index values. And now it is working.</p>
<pre><code class="lang-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 &lt; 0] = 0
    idx[idx &gt; max_value] = 0
    return torch.take(params, idx)
</code></pre>

ptrblck: <p>Thanks for debugging this issue!</p>
<aside class="quote no-group quote-modified" data-username="hzxie" data-post="2" data-topic="37502">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hzxie/48/4502_2.png" class="avatar"> hzxie:</div>
<blockquote>
<p>According to TF’s API:</p>
<blockquote>
<p>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.</p>
</blockquote>
</blockquote>
</aside>
<p>That seems like a weird strategy to me.</p>

hzxie: <p>I found that the results are not the same as TF when batch size is larger than 1.</p>
<p>For example, the size of <code>sample_idx</code> and <code>sample_grid</code> are <code>[57802752, 4]</code> and <code>[6, 32, 32, 32]</code>, respectively.</p>
<pre><code class="lang-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')
</code></pre>
<p>However, in TF:</p>
<pre><code class="lang-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
</code></pre>
<p>You can download the <code>sample_grid.npy</code> and <code>sample_idx.npy</code> from <a href="https://file.io/NNKgZF" rel="nofollow noopener">here</a>.</p>
<p>Who can help me out <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=6" title=":frowning:" class="emoji" alt=":frowning:"></p>

hzxie: <p>I wrote a small program to test the behavior of <code>TF</code> and the self-implemented <code>gather_nd</code> in PyTorch.</p>
<pre><code class="lang-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 &lt; 0] = 0
    idx[idx &gt; 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)
</code></pre>
<p>And the output is shown as following:</p>
<pre><code class="lang-auto">TF [ 0  5 25  0  0]
Torch tensor([  0,   5,  25,   0, 100])
</code></pre>
<p>It seem that there are something wrong with the self-implemented one.</p>

hzxie: <p>I preprocessed the indices before <code>torch.take</code> and make sure that the index value in each dimension is correct.</p>
<p>The modified <code>gather_nd</code> is listed as follows:</p>
<pre><code class="lang-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 &lt; 0] = 0
    for idx, ps in enumerate(params_size):
        indices[indices[:, idx] &gt;= 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)
</code></pre>

qingping: <p>Hi, I modified <a class="mention" href="/u/hzxie">@hzxie</a> ‘s code as below:<br>
It support the case that <code>indices.size(1)</code> is not equal to <code>len(param_size)</code></p>
<pre><code class="lang-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) &gt;= 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]
</code></pre>
<p>Hope that code helps someone <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>

Ren_yf: <p>I am also quite confused with torch implementation of  <code>tf.gather_nd</code> at first, but I found <code>tf.gather_nd</code> seems like a variant of slicing. I have my implementation below and hope it will help someone <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=9" title=":smile:" class="emoji" alt=":smile:"></p>
<pre><code class="lang-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 &lt;= n)"""
  assert type(indices) == torch.Tensor
  return params[indices.transpose(0,1).long().numpy().tolist()]
</code></pre>
<h2>testing</h2>
<p><a href="https://www.tensorflow.org/versions/r1.15/api_docs/python/tf/gather_nd" rel="nofollow noopener">Tensorflow samples</a></p>
<pre><code class="lang-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']
</code></pre>
<p>for torch</p>
<pre><code class="lang-python">&gt;&gt;&gt; params = torch.rand(2,2,2)
&gt;&gt;&gt; params
tensor([[[0.4685, 0.2514],
         [0.0624, 0.0797]],

        [[0.4989, 0.1414],
         [0.6970, 0.6825]]])
&gt;&gt;&gt; indices = torch.Tensor([[1]])
&gt;&gt;&gt; gather_nd(params, indices)
tensor([[[0.4989, 0.1414],
         [0.6970, 0.6825]]])
&gt;&gt;&gt; indices = torch.Tensor([[0, 1], [1, 0]])
&gt;&gt;&gt; gather_nd(params, indices)
tensor([[0.0624, 0.0797],
        [0.4989, 0.1414]])
&gt;&gt;&gt; indices = torch.Tensor([[0, 0, 1], [1, 0, 1]])
&gt;&gt;&gt; gather_nd(params, indices)
tensor([0.2514, 0.1414])
</code></pre>

Joris_Mollinga: <p>I created a version that works on 4D indices, 3D params and returns a 4D output.</p>
<p>Example:</p>
<pre><code class="lang-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]
</code></pre>
<p>This implementation is <strong>very</strong> slow and ugly. But maybe someone can think of a clever trick to speed up the <code>for</code> loop?</p>
<pre><code class="lang-auto">def torch_gather_nd(params: torch.Tensor,
                    indices: torch.Tensor) -&gt; 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] &lt;= 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
</code></pre>