Usage of torch.scatter() for multi-dimensional value

Yuxuan_Xue: <p>Hi!</p>
<p>i have a question regarding the usage of the torch.scatter() function.</p>
<p>I want to construct a weights matrix weights (# [B, N, V]. B is batch size, N is number of points and V is the number of features for each point. )</p>
<p>Let’s say i have two tensors</p>
<pre><code class="lang-auto">a = # shape [B, N, k], where B is batch size, N is number of points, k is the index number within [0,V] to select feature.
b = # shape [B, N, k], where B is batch size, N is number of points, k stores here the weights for selected feature.
</code></pre>
<p>I tried to use function torch.scatter():<br>
weights.scatter_(index=a, dim=2, value=some_fix_value). By this operation i can only set one fixed value, but not the whole value tensor b, which contains all information at those location.</p>
<p>Can someone gives me a hint on how to do this properly?</p>

ptrblck: <p>Direct indexing should work if I understand your use case correctly:</p>
<pre><code class="lang-python">B, N, V = 2, 3, 4

a = torch.randn(B, N, V)
k = 2
idx = torch.randint(0, V, (B, N, k))

b = torch.zeros(B, N, V)

r = torch.arange(b.size(0))[:, None, None]
c = torch.arange(b.size(1))[None, :, None]

b[r, c, idx] = a[r, c, idx]
print(idx)
# tensor([[[0, 1],
#          [2, 1],
#          [1, 2]],

#         [[0, 3],
#          [0, 3],
#          [3, 1]]])
print(a)
# tensor([[[ 1.0250, -1.8356,  0.4314,  2.1630],
#          [-0.8884, -0.2196, -1.5033,  0.8229],
#          [-0.1390, -0.9114, -0.2310, -0.7310]],

#         [[ 0.7332, -1.5779, -0.4527, -1.6785],
#          [ 0.6614, -0.0094, -0.1890,  0.3890],
#          [-0.4206, -1.1668, -0.1563, -0.6945]]])
print(b)
# tensor([[[ 1.0250, -1.8356,  0.0000,  0.0000],
#          [ 0.0000, -0.2196, -1.5033,  0.0000],
#          [ 0.0000, -0.9114, -0.2310,  0.0000]],

#         [[ 0.7332,  0.0000,  0.0000, -1.6785],
#          [ 0.6614,  0.0000,  0.0000,  0.3890],
#          [ 0.0000, -1.1668,  0.0000, -0.6945]]])
</code></pre>

Yuxuan_Xue: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>,</p>
<p>thank you! It is generally what i mean, but the matrix a in my case is # [B, N, k] instead of # [B, N, V]. I only know the value at k-th location, but i want to assign 0 to all other unselected locations. In other words, i am turning matrix a from #[B, N, k] to #[B, N, V], with original value in k-index and 0 in other location. Do you have an idea?</p>

ptrblck: <p>Isn’t this exactly what my code is doing?<br>
It selects values from a matrix <code>a</code> in the shape <code>[B, N, V]</code> using indices in <code>[B, N, k]</code> and assigns these to <code>b</code>, which is initialized with zeros for all other values.<br>
If not, could you post a slow reference implementation, please?</p>

Yuxuan_Xue: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Thank you, you are right.</p>

yunjiangster: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Thanks for the concise answer. I tried to generalize your solution to broadcast the last dimension, but somehow that didn’t work. Any suggestion on how to do it properly?</p>
<pre><code class="lang-auto">import torch
B, N, V = 2, 3, 4

a = torch.randn(B, N, V)
k = 2
idx = torch.randint(0, N, (B, k))

b = torch.zeros(B, N, V)
b[:,idx,:] = a[:,idx,:]
</code></pre>
<pre><code class="lang-auto">print(a)
tensor([[[-0.2632,  0.0441, -1.0859,  0.7385],
         [-0.6510, -1.0140, -1.0576,  0.6398],
         [ 0.0409,  0.2688, -1.1641,  0.8304]],

        [[ 0.2711, -1.5461,  0.4596, -0.3386],
         [ 0.7731, -1.2105, -0.3453,  0.7746],
         [-2.0150,  0.2547, -1.4397,  0.2060]]])

print(b)
tensor([[[-0.2632,  0.0441, -1.0859,  0.7385],
         [-0.6510, -1.0140, -1.0576,  0.6398],
         [ 0.0409,  0.2688, -1.1641,  0.8304]],

        [[ 0.2711, -1.5461,  0.4596, -0.3386],
         [ 0.7731, -1.2105, -0.3453,  0.7746],
         [-2.0150,  0.2547, -1.4397,  0.2060]]])
</code></pre>

yunjiangster: <p>Never mind, figured it out:</p>
<pre><code class="lang-auto">import torch
B, N, V = 2, 3, 4

a = torch.randn(B, N, V)
k = 2
idx = torch.randint(0, N, (B, k))

b = torch.zeros(B, N, V)
r = torch.arange(b.size(0))[:, None]

b[r, idx] = a[r, idx]
</code></pre>