Assign a tensor to multiple slices

Question

Let

a = tensor([[0, 0, 0, 0],
            [0, 0, 0, 0],
            [0, 0, 0, 0]])
b = torch.tensor([1, 2])
c = tensor([[1, 2, 0, 0],
            [0, 1, 2, 0],
            [0, 0, 1, 2]])

Is there a way to obtain c by assigning b to slices of a without any loops? That is, a[indices] = b for some indices or something similar?

Answer 1

You can use scatter method in pytorch.

a = torch.tensor([[0, 0, 0, 0],
                 [0, 0, 0, 0],
                 [0, 0, 0, 0]])

b = torch.tensor([1, 2])

index = torch.tensor([[0,1],[1,2],[2,3]])

a.scatter_(1, index, b.view(-1,2).repeat(3,1))
# tensor([[1, 2, 0, 0],
#         [0, 1, 2, 0],
#         [0, 0, 1, 2]])