How to assign to specific columns in each row?

1

Suppose I have a list of indices: ind = torch.tensor([0, 2, 1, 3]), and a 4x4 tensor A.

I need to assign some values to the items A[i, ind[i], for i=1,4. For example:

for i in range(4):
    A[i, ind[i]] = r[i]

where r contains the target values in this example.

Given A, ind, r, Is there a way to vectorize this operation?

1 Like
2

This should work:

# Setup
ind = torch.tensor([0, 2, 1, 3])
A = torch.zeros(4, 4)
r = torch.randn(4)

# Your code
for i in range(4):
    A[i, ind[i]] = r[i]

# Alternative
B = torch.zeros(4, 4)
B[torch.arange(4), ind] = r
print((B == A).all())
> tensor(True)
2 Likes