How to implement this with pytorch function (subtitude for For loops)?

ljmzlh: <p>How to implement this (essentially calculating average of vectors)</p>
<pre><code class="lang-python">batch_size, d= 8, 128
n, N= 100, 10

feature=torch.rand(batch_size, n, d)
belong=torch.randint(0, N, (batch_size,n))
output=torch.zeros(batch_size, N, d)

for i in range(batch_size):
    for j in range(N):
        output[i][j]=feature[i][belong[i]==j].mean(0)
</code></pre>
<p>with high-performance, CUDA-friendly and autograd-friendly pytorch function?<br>
Grateful to your help <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
P.S. I tried torch.index_add and it performs badly for CUDA</p>

KFrank: <p>Hi ljmzlh!</p>
<aside class="quote no-group" data-username="ljmzlh" data-post="1" data-topic="134514" 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/ljmzlh/48/41706_2.png" class="avatar"> ljmzlh:</div>
<blockquote>
<p>How to implement this</p>
</blockquote>
</aside>
<p>The core issue is that <code>feature[i][belong[i]==j]</code> has a different<br>
shape for different values of <code>i</code> and <code>j</code>.  Therefore if you try to build a<br>
“tensor” whose “slices” are given by the above expression, you will<br>
end up with a “ragged tensor” (that is, a tensor whose slices have<br>
differing shapes).  And pytorch doesn’t support ragged tensors.</p>
<p>Once you use such boolean “indexing” (the cause of the differing<br>
shapes), there is no way for you to complete your calculation using<br>
pure tensor operations.</p>
<p>Instead, you should convert the boolean expression <code>belong[i]==j</code><br>
to a <code>0</code>, <code>1</code> numerical mask and use element-wise multiplication to mask<br>
the desired elements of <code>feature[i]</code>.  (The boolean tensor will be<br>
automatically cast to a numerical tensor in the context we use it, so<br>
we don’t have to do it explicitly.)</p>
<p>Also, because you are calculating <code>mean()</code> rather than <code>sum()</code>, it proves<br>
convenient to use a <code>0</code>, <code>1 / nTrue</code> mask, rather than a <code>0</code>, <code>1</code> mask.</p>
<p>Lastly, there is no complication in using tensor operations to perform<br>
the calculation on all batch elements (all slices along the <code>batch_size</code><br>
dimension) at the same time.</p>
<p>Here is an illustrative script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.manual_seed (2021)

batch_size, d= 8, 128
n, N= 100, 10

feature=torch.rand(batch_size, n, d)
belong=torch.randint(0, N, (batch_size,n))

# boolean "indexing" with loop
output=torch.zeros(batch_size, N, d)

for i in range(batch_size):
    for j in range(N):
        output[i][j]=feature[i][belong[i]==j].mean(0)

# loop-free tensor method
mask = belong.unsqueeze (-1) == torch.arange (N)
mask = mask / mask.sum (1).unsqueeze (1)   # to compute mean() rather than sum()
outputB = torch.einsum ('ijk, ijl -&gt; ikl', mask, feature)

print ('outputB.allclose (output) =', outputB.allclose (output))
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">1.9.0
outputB.allclose (output) = True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

ljmzlh: <p>That’s helpful, thank you!</p>