PyTorch equivalent of tf.gather

username: Superklez
cooked: <p>So I have an input tensor with shape <code>[16, 1, 125, 256]</code> and a selector tensor with shape <code>[124, 2]</code>. Is there any PyTorch equivalent of <code>tf.gather(input, selector, axis=2)</code>? To begin with, how does Tensorflow handle both tensors despite not having the same number of dimensions (unlike <code>torch.gather</code> wherein they must be equal)? Furthermore, what must <code>dim</code> in <code>torch.gather</code> be to be similar to <code>axis=2</code> in <code>tf.gather</code>.</p>
<p>For more context, I’m trying to make PyTorch version of <a href="https://github.com/tensorflow/tensorflow/blob/r1.12/tensorflow/contrib/signal/python/ops/shape_ops.py" rel="noopener nofollow ugc">this</a>, which defines a <code>frame</code> function that expands <code>signal</code>'s <code>axis</code> dimension into frames of <code>frame_length</code>.</p>

username: albanD
cooked: <p>Could you explain what <code>tf.gather(input, selector, axis=2)</code> is doing when the there is no matching dimension between input and selector? Maybe with a small example of what you are doing?</p>

username: Superklez
cooked: <p>It gathers slices of select indices (<code>selector</code>) from <code>input</code> with respect to the specified <code>axis</code>. A workaround I found is <code>input[:, :, selector, :]</code> which returns the same output as <code>tf.gather(input, selector, axis=2)</code>. However, how do I do the same thing if the shape of <code>input</code> is not known? Here are some more examples:</p>
<p><code>tf.gather(input, selector, axis=0)</code> is the same as <code>input[selector, ...]</code><br>
<code>tf.gather(input, selector, axis=1)</code> is the same as <code>input[:, selector, ...]</code><br>
<code>tf.gather(input, selector, axis=2)</code> is the same as <code>input[..., selector, :]</code><br>
<code>tf.gather(input, selector, axis=3)</code> is the same as <code>input[..., selector]</code></p>

username: albanD
cooked: <p>It’s still not 100% clear to me what this does <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"> . But I guess that would work:</p>
<pre><code class="lang-auto">dim = 2
new_size = inp.size()[:dim] + selector.size() + inp.size()[dim+1:]
out = inp.index_select(selector.view(-1), dim=dim)
out.view(new_size)
</code></pre>

username: Superklez
cooked: <p>It works. Thanks for this! One last question: is there a way to simultaneously slice along all dimensions of a tensor. Say, for example, I have a tensor <code>inp</code> with shape <code>[16, 1, 125, 256]</code>. I then have two lists (this is a simple example)</p>
<p><code>begin = [0, 0, 0, 0]</code><br>
<code>end = [16, 1, 125, 256]</code></p>
<p>that specify the beginning and end indices to slice, respectively. From this, I perform slicing as</p>
<p><code>out = inp[begin[0]:end[0], begin[1]:end[1], begin[2]:end[2], begin[3]:end[3]]</code></p>
<p>How do I do this with any tensor with a variable number of dimensions?</p>

username: albanD
cooked: <p>The way I would do this is with a serie of narrows (keep in mind that narrow never copies the memory so it is very fast):</p>
<pre><code class="lang-python">out = inp
for dim, (b, e) in enumerate(zip(begin, end)):
   out = out.narrow(dim, b, e-b)</code></pre>

username: Superklez
cooked: <p>These work as expected. Thanks again for your help!</p>

username: ecolss
cooked: <p>Hey <a class="mention" href="/u/alband">@albanD</a>, I just came across <code>index_select</code>, looks like it’s a simpler version of <code>torch.gather</code>?<br>
For 2d matrix, <code>index_select</code> can only select same columns for all rows, but <code>gather</code> can select different columns for different rows.</p>
<p>I’m actually trying to use <code>gather</code>, but it failed with some errors, (see <a href="https://discuss.pytorch.org/t/torch-gather-doesnt-work-during-backward/151693" class="inline-onebox">Torch.gather() doesn't work during backward()</a>)</p>