How to understand the code "cc[bb] += aa" in pytorch?

Viewed 74
import torch
aa=torch.tensor([[1,2,3],[4,5,6]]).T
bb=torch.tensor([0,1,1]).T
cc = torch.zeros(2, 2)
cc[bb] += aa

Then the result is cc=tensor([[1., 4.],[3., 6.]]), why?

1 Answers

Let us reason with pseudo code:

>>> aa = [[1,2,3],
          [4,5,6]].T
>>> aa
[[1,4],
 [2,5],
 [3,6]]

>>> bb = [0,1,1].T
>>> b
[0,
 1,
 1]

>>> cc = zeros(2, 2)
>>> cc
[[0,0],
 [0,0]]

The next instruction is an assignment which consists in first indexing cc with bb values. Here we are picking entire rows from cc using indices in bb. Since there are three rows in bb, the resulting tensor will consist of cc[bb[0]], cc[bb[1]], and cc[bb[2]] but bb[1] and bb[2] are equal which means it comes down to cc[0] and cc[1].

The right-hand side operand is aa and consists of three rows: [1,4], [2,5], and [3,6]. This means that the final operation performed will be equivalent to (row-wise):

cc[0] += [1,4]
cc[1] += [3,6]

Since cc is initialized with zero values, we can sum this up to:

>>> cc[0] = [1,4]
>>> cc[1] = [3,6]

This means that:

>>> cc
[[1,4],
 [3,6]]
Related