Batched index_select

Vlad_Zhukov: <p>How can I use torch.index_select for 3D tensors?<br>
I need some function which will apply usual index_select per batch.<br>
F.e. I want to be able to do something like this:</p>
<pre><code class="lang-auto">A = torch.FloatTensor([[[1, 2, 3], [1, 2, 3]], [[4, 5, 6], [4, 5, 6]], [[7, 8, 9], [7,8, 9]]])
ind = torch.LongTensor([[0, 1], [1, 2], [0, 2]])
torch.index_select(A, 2, ind)
</code></pre>
<p>What is the most effective and simple way to do this?<br>
t.y.</p>

richard: <p>I’m not sure if this is the best way to do it, but this works:</p>
<pre><code class="lang-auto">torch.cat([ torch.index_select(a, 1, i).unsqueeze(0) for a, i in zip(A, ind) ])
</code></pre>
<p>Essentially what this does is apply the regular index_select to each batch-element of A and ind and then concatenate the results together.</p>

Vlad_Zhukov: <p>yeah, it’s pretty straightforward. I somehow thought that we can do it without “for loops”…I’m probably overextending <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:">  thank you.</p>

wasiahmad: <p>Is there any solution of batched <code>index_select</code> without using a loop?</p>

Vlad_Zhukov: <p>This is a bit out of context. But I hope one can figure out what is going here</p>
<pre><code class="lang-auto">def get_selected_matrices(probs, references, dim=1, sanity_check=False):
    """
    batched index select
    probs - is a matrix
    references - is index
    dim - is dimention of element of the batch
    """
    #return torch.cat([torch.index_select(a, dim, Variable(LongTensor(i))).unsqueeze(0)\
    #                        for a, i in zip(probs, references)])

    batch_size, seq_len, vocab_size = probs.size()
    references = torch.from_numpy(np.array(references)).long() # batch_size x seq_len
    vocab_extension = torch.arange(0, batch_size).long() * vocab_size # batch_size
    if torch.cuda.is_available():
        references = references.cuda()
        vocab_extension = vocab_extension.cuda()
    references_extended_vocab = (references + vocab_extension.unsqueeze(-1)).view(-1) # batch_size * seq_len
    probs_extended_vocab = torch.transpose(probs, 0, 1).contiguous().view(seq_len, -1) # seq_len x batch_size * vocab_size
    probs_reduced_extended_vocab = torch.index_select(
        probs_extended_vocab, dim, Variable(references_extended_vocab)
    ) # seq_len x batch_size * seq_len
    probs_reduced_vocab = torch.transpose(
        probs_reduced_extended_vocab.view(seq_len, batch_size, seq_len), 0, 1
    ) # batch_size x seq_len x seq_len
    if sanity_check:
        probs_reduced_vocab_loop = torch.cat(
            [torch.index_select(a, dim, Variable(LongTensor(i))).unsqueeze(0) for a, i in zip(probs, references)]
        )
        if not torch.equal(probs_reduced_vocab, probs_reduced_vocab_loop):
            raise AssertionError('got wrong probs with reduced vocab')
            print(probs_reduced_vocab)
            print(probs_reduced_vocab_loop)
    return probs_reduced_vocab
</code></pre>

bobchennan: <p>This is just converting two dimension index into one dimension.</p>

jnhwkim: <p>Is this what you want? <a href="http://pytorch.org/docs/master/torch.html#torch.gather" rel="nofollow noopener"><code>gather</code></a> might be useful for the task</p>
<pre><code class="lang-auto"># Batched index_select
def batched_index_select(t, dim, inds):
    dummy = inds.unsqueeze(2).expand(inds.size(0), inds.size(1), t.size(2))
    out = t.gather(dim, dummy) # b x e x f
    return out
</code></pre>

marcman411: <p>Piggy-backing on <a class="mention" href="/u/jnhwkim">@jnhwkim</a>’s idea of using gather, here is a function that should mimic <code>index_select</code> for arbitrary dimensional inputs. It just doesn’t work with negative indexing.</p>
<pre><code class="lang-auto"># input: B x * x ... x *
# dim: 0 &lt;= scalar
# index: M
def batched_index_select(input, dim, index):
	views = [1 if i != dim else -1 for i in xrange(len(input.shape))]
	expanse = list(input.shape)
	expanse[dim] = -1
	index = index.view(views).expand(expanse)
	return torch.gather(input, dim, index)
</code></pre>
<p>And here’s a variation if you have different indexes for each batch:</p>
<pre><code class="lang-auto"># input: B x * x ... x *
# dim: 0 &lt; scalar
# index: B x M
def batched_index_select(input, dim, index):
	views = [input.shape[0]] + \
		[1 if i != dim else -1 for i in xrange(1, len(input.shape))]
	expanse = list(input.shape)
	expanse[0] = -1
	expanse[dim] = -1
	index = index.view(views).expand(expanse)
	return torch.gather(input, dim, index)
</code></pre>

ratishsp: <aside class="quote no-group" data-username="jnhwkim" data-post="7" data-topic="9115">
<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/jnhwkim/48/1036_2.png" class="avatar"> jnhwkim:</div>
<blockquote>
<p>def batched_index_select(t, dim, inds): dummy = inds.unsqueeze(2).expand(inds.size(0), inds.size(1), t.size(2)) out = t.gather(dim, dummy) # b x e x f return out</p>
</blockquote>
</aside>
<p>Hi <a class="mention" href="/u/jnhwkim">@jnhwkim</a>,<br>
The expected output is<br>
(0 ,.,.) =<br>
1  2<br>
1  2</p>
<p>(1 ,.,.) =<br>
5  6<br>
5  6</p>
<p>(2 ,.,.) =<br>
7  9<br>
7  9<br>
[torch.FloatTensor of size 3x2x2]</p>
<p>Running your version of batched_index_select results in a different output<br>
(0 ,.,.) =<br>
1  1  1<br>
2  2  2</p>
<p>(1 ,.,.) =<br>
5  5  5<br>
6  6  6</p>
<p>(2 ,.,.) =<br>
7  7  7<br>
9  9  9<br>
[torch.FloatTensor of size 3x2x3]</p>

Zhaoyi-Yan: <p>Maybe a slightly better version is here.</p>
<pre><code class="lang-auto">import torch
# input: B x * x ... x *
# dim: 0 &lt;= scalar
# index: M
def batched_index_select(input, dim, index):
    views = [1 if i != dim else -1 for i in range(len(input.shape))]
    expanse = list(input.shape)
    expanse[dim] = -1
    index = index.view(views).expand(expanse)
    # making the first dim of output be B
    return torch.cat(torch.chunk(torch.gather(input, dim, index), chunks=index.shape[0], dim=dim), dim=0)
</code></pre>

dashesy: <aside class="quote no-group" data-username="marcman411" data-post="8" data-topic="9115">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/7cd45c/48.png" class="avatar"> marcman411:</div>
<blockquote>
<p>different indexes for each batch</p>
</blockquote>
</aside>
<p>That is exactly what I wanted, thanks. Since I did not want to rely on the exact batch number I used this variant</p>
<pre data-code-wrap="python"><code class="lang-python">def batched_index_select(input, dim, index):
    for ii in range(1, len(input.shape)):
        if ii != dim:
            index = index.unsqueeze(ii)
    expanse = list(input.shape)
    expanse[0] = -1
    expanse[dim] = -1
    index = index.expand(expanse)
    return torch.gather(input, dim, index)
</code></pre>

spadel: <p>I have a similar task, so maybe someone can help me!</p>
<p>So what I have are two tensors:</p>
<ul>
<li>an indices tensor <code>indices</code> with shape <code>(2, 5, 2)</code>, where the last dimensions corresponds to indices in x and y dimension</li>
<li>a “value tensor” <code>value</code> with shape <code>(2, 5, 2, 16, 16)</code>, where I want the last two dimensions to be selected with x and y indices</li>
</ul>
<p>To be more concrete, the indices are between 0 and 15 and I want to get an output:</p>
<pre><code>out = value[:, :, :, x_indices, y_indices]
</code></pre>
<p>For example the indices for the first item are <code>i_y = indices[0, 0, 0]; i_x = indices[0, 0, 1]</code>. The corresponding output should then be <code>out[0, 0, :, i_y, i_x]</code>. Can anybody help me here? Thanks a lot!</p>

eacousineau: <p>Not sure if this is comparable to <a href="https://discuss.pytorch.org/u/dashesy">dashesy</a>’s in terms of generality, but here’s a simplified case + tests for selecting vectors, e.g. via <code>argmin</code> or <code>multinomial</code>:</p>
<aside class="onebox githubgist" data-onebox-src="https://gist.github.com/EricCousineau-TRI/cc2dc27c7413ea8e5b4fd9675050b1c0">
  <header class="source">

      <a href="https://gist.github.com/EricCousineau-TRI/cc2dc27c7413ea8e5b4fd9675050b1c0" target="_blank" rel="noopener nofollow ugc">gist.github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://gist.github.com/EricCousineau-TRI/cc2dc27c7413ea8e5b4fd9675050b1c0" target="_blank" rel="noopener nofollow ugc">https://gist.github.com/EricCousineau-TRI/cc2dc27c7413ea8e5b4fd9675050b1c0</a></h4>

  <h5>vector_gather.py</h5>
  <pre><code class="Python">import einops
import torch

def vector_gather(vectors, indices):
    """
    Gathers (batched) vectors according to indices.

    Arguments:
        vectors: Tensor[N, L, D]
        indices: Tensor[N, K] or Tensor[N]</code></pre>
    This file has been truncated. <a href="https://gist.github.com/EricCousineau-TRI/cc2dc27c7413ea8e5b4fd9675050b1c0" target="_blank" rel="noopener nofollow ugc">show original</a>
  <h5>vector_gather_test.py</h5>
  <pre><code class="Python">import unittest

import einops
import torch


class Test(unittest.TestCase):
    def test_vector_gather(self):
        vector = torch.tensor([
            [</code></pre>
    This file has been truncated. <a href="https://gist.github.com/EricCousineau-TRI/cc2dc27c7413ea8e5b4fd9675050b1c0" target="_blank" rel="noopener nofollow ugc">show original</a>

<p>
</p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


Zohaib: <p>Thanks, <a class="mention" href="/u/eacousineau">@eacousineau</a>. I was looking for this.</p>