Updating elements in-place in one tensor given specific locations at another tensor

1

Hi all!
I am stuck at one problem. Let say that that I have tensors A and B. Now, I want to update elements in A at locations given in tensor B.
For example, let say there are two tensors A and B with size 3 x 6, and 3 x 2.

A = [
    [1, 2, 3, 4, 5, 6],
    [2, 3, 4, 5, 6, 1],
    [3, 4, 5, 6, 1, 2],
]
B = [
    [0, 1],
    [0, 3],
    [1, 4],
]

I wanna do something like this:

A[B] += 100 # I want the operation to be vectorized.

So that I will get the following result:

A = [
    [101, 102, 3, 4, 5, 6],
    [102, 3, 4, 105, 6, 1],
    [3, 104, 5, 6, 101, 2],
]
2

Hi Nauryzbay!

You may use scatter_add_() (with pytorch tensors) for this:

>>> import torch
>>> torch.__version__
'1.9.0'
>>> A = torch.tensor ([
...     [1, 2, 3, 4, 5, 6],
...     [2, 3, 4, 5, 6, 1],
...     [3, 4, 5, 6, 1, 2],
... ])
>>> B = torch.tensor ([
...     [0, 1],
...     [0, 3],
...     [1, 4],
... ])
>>> A.scatter_add_ (1, B, torch.tensor ([100]).expand (B.shape))
tensor([[101, 102,   3,   4,   5,   6],
        [102,   3,   4, 105,   6,   1],
        [  3, 104,   5,   6, 101,   2]])
>>> A
tensor([[101, 102,   3,   4,   5,   6],
        [102,   3,   4, 105,   6,   1],
        [  3, 104,   5,   6, 101,   2]])

Best.

K. Frank

1 Like