Indexed copy along dimension

1

I want to copy tensor elements to another tensor at positions given by indices along given dimension.
Example:
A = [[3, 2, 1], [2, 4, 5], [6, 3, 2]]
B = [X, Y, Z]
Index = [0, 2, 1]
Dim = 1

Result should be: [[X, 2, 1], [2, 4, Y], [6, Z, 2]]
What’s the best way to achieve this?

Thanks!

2

Hi Umbriel!

You could use indexing with assignment:

>>> import torch
>>> torch.__version__
'1.10.2'
>>> A = torch.tensor ([[3, 2, 1], [2, 4, 5], [6, 3, 2]])
>>> B = torch.tensor ([100, 200, 300])
>>> Index = torch.tensor ([0, 2, 1])
>>> A[torch.arange (3), Index] = B
>>> A
tensor([[100,   2,   1],
        [  2,   4, 200],
        [  6, 300,   2]])

Best.

K. Frank

3

Thank you KFrank, it works!