Assign row data to a 3D tensor on indices

asabater: <p>Hi! I have a multidimensional Tensor (4 or more dimensions) and I want to assign data to it based on given indices.<br>
E.g. in the tensor <em>mem</em> I want to modify the 6 rows (3 rows of each dimension 0) with the indices <em>indices</em> and their final value must be the one given by the 6 rows on <em>values</em>.:</p>
<pre><code class="lang-auto">mem = torch.zeros(2,4,5,6)                         # size -&gt; (2,4,5,6)
indices = torch.Tensor([[[0, 0],[0, 1],[2, 1]],    # size -&gt; (2,3,2)
        [[1, 0],[0, 0],[2, 0]]]).long()
values = torch.arange(1,2*3*6+1).reshape(2,3,6)    # size -&gt; (2,3,6)
mem.index_add(0, indices, values)
</code></pre>
<p>However, I find the following error:<br>
<code>IndexError: index_add_(): Index is supposed to be a vector</code></p>
<p>How can I address this problem when using Tensor with 4 or more dimensions?<br>
Note that both indices and values share the first two dimension sizes, and the last dimension of the values correspond to the last dimension of the mem Tensor.<br>
Thank you in advance.</p>

ptrblck: <p>I don’t fully understand the example as the number of <code>indices</code> is smaller than the <code>values</code>.<br>
Could you write a (slow) approach using loops to see what the desired output would be?</p>

asabater: <p>I had a coding mistake, already modified in the code above.<br>
The solution using loops is as follows:</p>
<pre><code class="lang-auto">mem = torch.zeros(2,4,5,6)                         # size -&gt; (2,4,5,6)
indices = torch.Tensor([[[0, 0],[0, 1],[2, 1]],    # size -&gt; (2,3,2)
        [[1, 0],[0, 0],[2, 0]]]).long()
values = torch.arange(1,2*3*6+1).reshape(2,3,6)    # size -&gt; (2,3,6)

B = indices.shape[0]
for b in range(B):
    for v, (i,j) in enumerate(indices[b]):
        mem[b,i,j] = values[b][v]
</code></pre>
<p>The result should be as follows (only the first value of each vector is shown):</p>
<pre><code class="lang-auto">print(mem[...,0])
tensor([[[ 1.,  7.,  0.,  0.,  0.],
         [ 0.,  0.,  0.,  0.,  0.],
         [ 0., 13.,  0.,  0.,  0.],
         [ 0.,  0.,  0.,  0.,  0.]],

        [[25.,  0.,  0.,  0.,  0.],
         [19.,  0.,  0.,  0.,  0.],
         [31.,  0.,  0.,  0.,  0.],
         [ 0.,  0.,  0.,  0.,  0.]]])
</code></pre>

KFrank: <p>Hi Alberto!</p>
<aside class="quote no-group" data-username="asabater" data-post="3" data-topic="145897" 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/asabater/48/30506_2.png" class="avatar"> asabater:</div>
<blockquote>
<p>The result should be as follows (only the first value of each vector is shown):</p>
<pre><code class="lang-auto">print(mem[...,0])
tensor([[[ 1.,  7.,  0.,  0.,  0.],
         [ 0.,  0.,  0.,  0.,  0.],
         [ 0., 13.,  0.,  0.,  0.],
         [ 0.,  0.,  0.,  0.,  0.]],

        [[25.,  0.,  0.,  0.,  0.],
         [19.,  0.,  0.,  0.,  0.],
         [31.,  0.,  0.,  0.,  0.],
         [ 0.,  0.,  0.,  0.,  0.]]])
</code></pre>
</blockquote>
</aside>
<p>You may create a view into <code>mem</code> using pytorch indexing and then<br>
assign <code>values</code> into <code>mem</code> through the indexed view:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

mem = torch.zeros(2,4,5,6)                         # size -&gt; (2,4,5,6)
indices = torch.Tensor([[[0, 0],[0, 1],[2, 1]],    # size -&gt; (2,3,2)
                        [[1, 0],[0, 0],[2, 0]]]).long()
values = torch.arange(1,2*3*6+1).reshape(2,3,6)    # size -&gt; (2,3,6)

# set up the indices
ind0 = torch.arange (2).unsqueeze (-1).unsqueeze (-1).expand (2, 3, 6)
ind1 = indices[:, :, 0].unsqueeze (-1).expand (2, 3, 6)
ind2 = indices[:, :, 1].unsqueeze (-1).expand (2, 3, 6)
ind3 = torch.arange (6).expand (2, 3, 6)

# assign into the indexed view
mem[ind0, ind1, ind2, ind3] = values.float()

print ('mem[..., 0]:')
print (mem[..., 0])
</code></pre>
<p>Here is the result:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; exec (open ('./index_assign.py').read())
1.10.2
mem[..., 0]:
tensor([[[ 1.,  7.,  0.,  0.,  0.],
         [ 0.,  0.,  0.,  0.,  0.],
         [ 0., 13.,  0.,  0.,  0.],
         [ 0.,  0.,  0.,  0.,  0.]],

        [[25.,  0.,  0.,  0.,  0.],
         [19.,  0.,  0.,  0.,  0.],
         [31.,  0.,  0.,  0.,  0.],
         [ 0.,  0.,  0.,  0.,  0.]]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>