Assign row data to a 3D tensor on indices

1

Hi! I have a multidimensional Tensor (4 or more dimensions) and I want to assign data to it based on given indices.
E.g. in the tensor mem I want to modify the 6 rows (3 rows of each dimension 0) with the indices indices and their final value must be the one given by the 6 rows on values.:

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

However, I find the following error:
IndexError: index_add_(): Index is supposed to be a vector

How can I address this problem when using Tensor with 4 or more dimensions?
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.
Thank you in advance.

2

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

3

I had a coding mistake, already modified in the code above.
The solution using loops is as follows:

mem = torch.zeros(2,4,5,6)                         # size -> (2,4,5,6)
indices = torch.Tensor([[[0, 0],[0, 1],[2, 1]],    # size -> (2,3,2)
        [[1, 0],[0, 0],[2, 0]]]).long()
values = torch.arange(1,2*3*6+1).reshape(2,3,6)    # size -> (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]

The result should be as follows (only the first value of each vector is shown):

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.]]])
4

Hi Alberto!

You may create a view into mem using pytorch indexing and then
assign values into mem through the indexed view:

import torch
print (torch.__version__)

mem = torch.zeros(2,4,5,6)                         # size -> (2,4,5,6)
indices = torch.Tensor([[[0, 0],[0, 1],[2, 1]],    # size -> (2,3,2)
                        [[1, 0],[0, 0],[2, 0]]]).long()
values = torch.arange(1,2*3*6+1).reshape(2,3,6)    # size -> (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])

Here is the result:

>>> 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.]]])

Best.

K. Frank

1 Like