Efficiently expand a tensor using higher-dimensional mask

gcortes: <p>Let us assume I’m working with a 2D tensor <code>data</code> of shape <code>(4, 1)</code> whose contents are given by <code>[[1], [2], [3], [4]]</code>. Now let <code>expanded_data = torch.empty(4, 2, 1)</code> and let <code>indices</code> be a 2D tensor of shape <code>(4, 2)</code> whose contents are given by <code>[[2, 3], [0, 2], [0, 1], [1, 2]]</code>. I would like to use <code>indices</code> as a mask to “expand” the original <code>data</code> tensor by populating <code>expanded_data</code> such that <code>expanded_data[i] = data[indices[i]]</code>. In this example, the contents of <code>expanded_in</code> should look like</p>
<p><code>[[[3], [4]], [[1], [3]], [[1], [2]], [[2], [3]]]</code>.</p>
<p>My naive implementation would be as follows:</p>
<pre><code class="lang-auto">for i in range(expanded_data.shape[0]):
    for idx, j in enumerate(indices[i]):
        expanded_data[i][idx] = data[j]
</code></pre>
<p>Is there a more efficient way to do this? A solution would preferably refrain from using for-loops and would remain entirely within the GPU.</p>

KFrank: <p>Hi Gcortes!</p>
<aside class="quote no-group" data-username="gcortes" data-post="1" data-topic="158879" data-full="true">
<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/g/d2c977/48.png" class="avatar"> gcortes:</div>
<blockquote>
<p>Is there a more efficient way to do this?</p>
</blockquote>
</aside>
<p>You can do this without loops by using  tensor indexing (with <code>squeeze()</code><br>
and <code>unsqueeze()</code> to get the shapes the way you want them):</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.12.0'
&gt;&gt;&gt; data = torch.tensor ([[1], [2], [3], [4]])
&gt;&gt;&gt; indices = torch.tensor ([[2, 3], [0, 2], [0, 1], [1, 2]])
&gt;&gt;&gt; expanded_data = torch.empty(4, 2, 1)
&gt;&gt;&gt; for i in range(expanded_data.shape[0]):
...     for idx, j in enumerate(indices[i]):
...         expanded_data[i][idx] = data[j]
...
&gt;&gt;&gt; expanded_dataB = data.squeeze()[indices.unsqueeze (2)].float()
&gt;&gt;&gt; torch.equal (expanded_data, expanded_dataB)
True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

gcortes: <p><a class="mention" href="/u/kfrank">@KFrank</a> I’m getting the following error from your solution:</p>
<pre><code class="lang-auto">IndexError: too many indices for tensor of dimension 1
</code></pre>

KFrank: <p>Hi Gcortes!</p>
<aside class="quote no-group" data-username="gcortes" data-post="3" data-topic="158879" data-full="true">
<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/g/d2c977/48.png" class="avatar"> gcortes:</div>
<blockquote>
<pre><code class="lang-auto">IndexError: too many indices for tensor of dimension 1
</code></pre>
</blockquote>
</aside>
<p>Could you post a minimal but complete and runnable script that reproduces<br>
your issue?</p>
<p>None of the tensors you mentioned in your original post are of dimension 1,<br>
and the script whose output I posted works.  You might try checking that the<br>
tensors in your failed version have dimensions consistent with those in your<br>
original post or in my example.</p>
<p>Best.</p>
<p>K. Frank</p>