Index select with potentially multiple indices

AfonsoSalgadoSousa: <p>I have a tensor <code>A</code>of size <code>B x 4</code> and another tensor <code>B</code> of size <code>H x N</code>. <code>A</code> has integer values and many are -1 like:</p>
<pre><code>tensor([[ 2.2640e+03, -1.0000e+00, -1.0000e+00, -1.0000e+00],
        [ 3.5400e+02, -1.0000e+00, -1.0000e+00, -1.0000e+00],
        [ 6.2700e+02, 6.2900e+02, -1.0000e+00, -1.0000e+00]
</code></pre>
<p>I want a new tensor of size <code>B x H</code> with values from <code>B</code> according to the indices in <code>A</code>. For example, for the first entry, I have the tensor [2260, -1, -1, -1] and I want the tensor of size <code>H,</code> in position <code>N</code> (i.e., tensor2[:, 2260]). There might be entries from <code>A</code> where there are more than one number higher than -1, in which case I would like to get the respective tensors in the indices and then average them.</p>
<p><code>torch.index_select(b, 0, a[:, 0])</code> almost gives the correct answer. Just need a way to handle the case where more than 1 integer is given, like in the third row of the above example.</p>
<p>Thanks for any help you can provide.</p>

KFrank: <p>Hi Afonso!</p>
<aside class="quote no-group quote-modified" data-username="AfonsoSalgadoSousa" data-post="1" data-topic="164255" data-full="true">
<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/afonsosalgadosousa/48/13258_2.png" class="avatar"> AfonsoSalgadoSousa:</div>
<blockquote>
<p>I want a new tensor of size <code>B x H</code> with values from <code>B</code> according to the indices in <code>A</code>.<br>
…<br>
I would like to get the respective tensors in the indices and then average them.</p>
</blockquote>
</aside>
<p>Because you want to average multiple entries of <code>B</code>, trying to use<br>
<code>index_select()</code> won’t fit your use case.</p>
<p>Use <code>A</code> to pluck the desired values out of <code>B</code> with pytorch tensor indexing<br>
and then compute the average “manually.”</p>
<p>Here is an illustration of such a scheme with a version of <code>B</code> that contains<br>
smaller index values:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt;
&gt;&gt;&gt; torch.__version__
'1.12.0'
&gt;&gt;&gt;
&gt;&gt;&gt; A = torch.tensor ([
...     [ 2, -1, -1, -1],
...     [ 3, -1, -1, -1],
...     [ 7,  6, -1, -1]
... ])
&gt;&gt;&gt;
&gt;&gt;&gt; B = torch.arange (50.0).view (5, 10) + 1
&gt;&gt;&gt; B
tensor([[ 1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10.],
        [11., 12., 13., 14., 15., 16., 17., 18., 19., 20.],
        [21., 22., 23., 24., 25., 26., 27., 28., 29., 30.],
        [31., 32., 33., 34., 35., 36., 37., 38., 39., 40.],
        [41., 42., 43., 44., 45., 46., 47., 48., 49., 50.]])
&gt;&gt;&gt;
&gt;&gt;&gt; Amask = A &gt; -1
&gt;&gt;&gt; Aind = torch.max (A, torch.zeros (1, dtype = torch.long))
&gt;&gt;&gt; result = ((B[:, Aind] * Amask).sum (dim = -1) / Amask.sum (dim = -1)).T
&gt;&gt;&gt;
&gt;&gt;&gt; result
tensor([[ 3.0000, 13.0000, 23.0000, 33.0000, 43.0000],
        [ 4.0000, 14.0000, 24.0000, 34.0000, 44.0000],
        [ 7.5000, 17.5000, 27.5000, 37.5000, 47.5000]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>