How to use Tf.gather_nd in pytorch

username: bdqnghi
cooked: <p>I have 2 tensors that have shapes like this:</p>
<p>(15, 720, 30)<br>
(15, 720, 19, 2)</p>
<p>I want to gather the tensors into one tensor with the shape like this : (15, 720, 19, 30), how can I do this?<br>
From the other topic: <a href="https://discuss.pytorch.org/t/how-to-do-the-tf-gather-nd-in-pytorch/6445/3">How to do the tf.gather_nd in pytorch?</a>, i see that there is a way but it’s tricky for me to understand how to use this.</p>

username: albanD
cooked: <p>Hi,</p>
<p>What should your output matrix contain? I fail to see what is the operation that you want to do here.</p>

username: bdqnghi
cooked: <p>Hi,</p>
<p>Thanks for your reply.</p>
<p>If you look at this repo at this line: <a href="https://github.com/bdqnghi/bi-tbcnn/blob/master/bi-tbcnn/bi-tbcnn/network.py#L155" rel="nofollow noopener">https://github.com/bdqnghi/bi-tbcnn/blob/master/bi-tbcnn/bi-tbcnn/network.py#L155</a>, which performs a bunch of complicated matrix operations on the tree structure.</p>
<p>These lines do the job:</p>
<pre><code class="lang-auto">vector_lookup = tf.concat([zero_vecs, nodes[:, 1:, :]], axis=1)

....
children = tf.concat([batch_indices, children], axis=3)
return tf.gather_nd(vector_lookup, children, name='children')
</code></pre>
<p>the shape of the vector lookup will be sth like: (15, 720, 30)<br>
and the shape of the children will be sth like: (15, 720, 19, 2)</p>
<p>and my goal is to do thing similar to the <code>return</code> line<br>
I’m trying to port this into Pytorch but seems quite tricky in some parts.</p>

username: enisberk
cooked: <p>I do not know tensorflow very well, so I am not sure if I got what you are asking correctly, can you check if thats what you are looking for:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; lookup=torch.ones((15, 720, 30))
&gt;&gt;&gt; children=torch.randint(0,15,(15, 720, 19, 2),dtype=torch.long)
&gt;&gt;&gt; lookup[children[:,:,:,0],children[:,:,:,1],:].shape
torch.Size([15, 720, 19, 30])
</code></pre>

username: bdqnghi
cooked: <p><a class="mention" href="/u/enisberk">@enisberk</a><br>
you’re the lifesaver !!!, thanks a lot</p>
<p>But I don’t understand how did you come up with the solution? Do you mind sharing how the code works?</p>
<p>The scenario can be described like this:<br>
<code>lookup=torch.ones((15, 720, 30))</code></p>
<p>means:  batch_size x num_nodes x feature_size<br>
where num nodes are the number of nodes in a tree and each node is represented by an embedding with size = feature_size.</p>
<p><code>children= (15, 720, 19, 2)</code><br>
means: batch_size x num_nodes x num_children x 2. I don’t really understand the meaning of 2 here.</p>
<p>Since every node in a tree is represented by an embedding, what the tensorflow code did is that they want to “merge” the 2 tensors into one, the dimension 4 means that each child in dimension 3 contains a corresponding embedding of its parent, make it more convenient for the later step.</p>
<p>In your code, i don’t really get how it works but it did the job…</p>

username: enisberk
cooked: <p>I am glad it helped.<br>
check that out first to understand gather_nd <a href="https://stackoverflow.com/a/42608263" rel="nofollow noopener">explanation of gather_nd from stackoverflow</a></p>
<p>So basically gather_nd needs set of indexes(children in our case) to use,<br>
In our case children have shape of <code>[15, 720, 19, 2]</code> so you can think it as, there are <code>15*720*19</code> index tuples with two elements, lets say one of those tuples is equal to (19,22), it corresponds to lookup[19,22,:] —&gt; this has shape of <code>[1,1,30]</code> but as I said there <code>15*720*19</code>  of those tuples, when you combine them all, you get<code>[15,720,19,30]</code>.</p>
<p>When you translate that to pytorch, all those index tuples corresponds to :<br>
<code>children[:,:,:,0],children[:,:,:,1]</code><br>
you need to use them as indexing elements to lookup, but you do not have an index for last dimension so you get all of them:<br>
<code>lookup[children[:,:,:,0],children[:,:,:,1],:]</code></p>
<p>I hope that makes it clearer.</p>

username: bdqnghi
cooked: <p>Thanks,</p>
<p>it’s clearer for me now <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=6" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: Nan_Meng
cooked: <p>That’s smart! And thanks your answer really helps me.</p>