Indexing a Tensor only partially with LongTensors

username: xeTaiz
cooked: <p>Hi there,<br>
I’m trying something along those lines:<br>
I have indices to a 3D volume (<code>b</code>), and <code>b.nonzero()</code> gives me (N, 3) indices to index with.<br>
Now I have another 3D volume with some more dimensions (<code>a</code>) and I would like to get something of the shape (BS, F, N) with N being the number of nonzero samples in <code>b</code>.</p>
<pre><code class="lang-auto">a = torch.randn(1,2, 10,10,10)
b = torch.randn(     10,10,10)
a[:, :, (b &gt; 0).nonzero(as_tuple=True)]
</code></pre>
<p>Leading to</p>
<pre><code class="lang-auto">TypeError: only integer tensors of a single element can be converted to an index
</code></pre>
<p>Obviously <code>b[(b &gt; 0).nonzero(as_tuple=True)]</code> works, but I can’t seem to index the first two dimensions differently (i.e. with something else but a LongTensor).</p>
<p>The only way I have in mind to do this would be a <code>torch.nn.functional.grid_sample()</code> with nearest mode after dividing the indices by the respective dimension lengths, but that seems unnecessary and slow.</p>
<p>Any suggestions on how I can achieve the above behavior?</p>

username: KFrank
cooked: <p>Hi Dominik!</p>
<aside class="quote no-group" data-username="xeTaiz" data-post="1" data-topic="162223" 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/xetaiz/48/9947_2.png" class="avatar"> xeTaiz:</div>
<blockquote>
<pre><code class="lang-auto">a = torch.randn(1,2, 10,10,10)
b = torch.randn(     10,10,10)
a[:, :, (b &gt; 0).nonzero(as_tuple=True)]
</code></pre>
</blockquote>
</aside>
<p>Use pytorch tensor indexing as follows:</p>
<p>Create five appropriately-structured index tensors – one each for each dimension<br>
of <code>a</code> – of shape <code>[BS, F, N]</code>, and use these to index into <code>a</code>.</p>
<p>Thus:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.12.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt;
&gt;&gt;&gt; a = torch.randn (1, 2, 10, 10, 10)
&gt;&gt;&gt; b = torch.randn (10, 10, 10)
&gt;&gt;&gt;
&gt;&gt;&gt; BS = a.size (0)
&gt;&gt;&gt; F = a.size (1)
&gt;&gt;&gt; N = (b &gt; 0).sum()
&gt;&gt;&gt;
&gt;&gt;&gt; ind_n0, ind_n1, ind_n2 = (b &gt; 0).nonzero (as_tuple = True)
&gt;&gt;&gt;
&gt;&gt;&gt; ind_bs = torch.arange (BS)[:, None, None].expand (BS, F, N)
&gt;&gt;&gt; ind_f = torch.arange (F)[None, :, None].expand (BS, F, N)
&gt;&gt;&gt; ind_n0 = ind_n0[None, None, :].expand (BS, F, N)
&gt;&gt;&gt; ind_n1 = ind_n1[None, None, :].expand (BS, F, N)
&gt;&gt;&gt; ind_n2 = ind_n2[None, None, :].expand (BS, F, N)
&gt;&gt;&gt;
&gt;&gt;&gt; result = a[ind_bs, ind_f, ind_n0, ind_n1, ind_n2]
&gt;&gt;&gt;
&gt;&gt;&gt; result.shape
torch.Size([1, 2, 511])
&gt;&gt;&gt;
&gt;&gt;&gt; # check result
&gt;&gt;&gt; torch.equal (result[0, 0], a[0, 0][(b &gt; 0).nonzero (as_tuple = True)])
True
&gt;&gt;&gt; torch.equal (result[0, 1], a[0, 1][(b &gt; 0).nonzero (as_tuple = True)])
True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>