How to assign a image like tensor with indexes

username: Ykk_L
cooked: <p>The target tensor is X (B, H, W, C). The index tensor is I (B, N, 2). And the source tensor is S (B, N, C).<br>
So how can I implement X[i, I[i,j,0], I[i,j,1]] = S[i, j] efficiently and parallelly?</p>
<p>I tried to use <code>for</code> loop to assign X. But it took too much time in <code>loss.backward()</code>.</p>

username: ptrblck
cooked: <p>This code might work:</p>
<pre><code class="lang-python">B, H, W, C = 2, 3, 3, 4
N = 5
x = torch.randn(B, H, W, C)
I = torch.randint(0, H, (B, N, 2))
S = torch.ones(B, N, C)

x[torch.arange(B).unsqueeze(1),
  I[torch.arange(B), :, 0],
  I[torch.arange(B), :, 1]] = S
</code></pre>
<p>Could you please double check it using your manual approach with the for loop?</p>
<p>PS: Please don’t create double posts, as other users might answer the same question <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>

username: Ykk_L
cooked: <p>It’s great.</p>
<p>I have tried to change the index matrix <code>I</code> to a new matrix with size <code>(B*N,3)</code> including the batch index.<br>
And this code works:</p>
<pre><code class="lang-auto">x[I[:,0], I[:,1], I[:2]] = S
</code></pre>
<p>However, your implementation is faster than mine.</p>
<p>And I have another question. If some indexes are repeating in <code>I</code>, which value in <code>S</code> would replace the value in <code>x</code>?  Is the latter one in <code>I</code>?</p>
<p>PS: This is the first time I post the problem on the forum. I’m sorry for my misoperation. Could you please delete another question? Thanks.</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="Ykk_L" data-post="3" data-topic="96686">
<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/ykk_l/48/28894_2.png" class="avatar"> Ykk_L:</div>
<blockquote>
<p>And I have another question. If some indexes are repeating in <code>I</code> , which value in <code>S</code> would replace the value in <code>x</code> ? Is the latter one in <code>I</code> ?</p>
</blockquote>
</aside>
<p>On the CPU it might be the last value. I think this behavior is undefined on the GPU.</p>
<aside class="quote no-group" data-username="Ykk_L" data-post="3" data-topic="96686">
<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/ykk_l/48/28894_2.png" class="avatar"> Ykk_L:</div>
<blockquote>
<p>PS: This is the first time I post the problem on the forum. I’m sorry for my misoperation. Could you please delete another question? Thanks.</p>
</blockquote>
</aside>
<p>Sure, we’ll delete the other post. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>

username: Ykk_L
cooked: <p>How to implement a similar task if the tensor is expended one more dimension?<br>
X (B, S, H, W, C)<br>
I (B, S, N, 2)<br>
S (B, S, N, C)<br>
The objective is X[i, j, I[i, j, k, 0], I[i, j, k, 1]] = S[i, j, k].</p>
<p>Thanks.</p>

username: Ykk_L
cooked: <p>I think this method could work</p>
<pre><code class="lang-auto">x[torch.arange(B).unsqueeze(1).unsqueeze(1),
  torch.arange(S).unsqueeze(1),
  I[torch.arange(B).unsqueeze(1), torch.arange(S), :, 0],
  I[torch.arange(B).unsqueeze(1), torch.arange(S), :, 1]] = S
</code></pre>