Hello Jet!



 jetcai1900:

I would like to use B to update A,
A[C[0]] += B[0]
A[C[1]] += B[1]
So far, I am using for loop to achieve this, which is very slow.
Any suggestions on how to get rid of for loop or speed up the implementation.


I am not aware of any way to do this in pytorch without a loop.
Furthermore, I doubt that there is such a way.  My reasoning runs
as follows:
The purpose of using “tensor” operations (whatever those are) is
to be able to do things “in parallel” (whatever that may mean).  But,
as you describe your desired calculation, it is inherently sequential.
Also, I don’t think it is possible to describe a useful semantics for
this type of operation that isn’t inherently sequential.
If all of the indices in your index tensor C were guaranteed to be
distinct, then one could parallelize the operations.  But your specific
example, C = [0, 0], provides a good counter-example.
Noting that C[0] = C[1] = 0, we can rewrite:
A[C[0]] += B[0]
A[C[1]] += B[1]

as:
A0tmpa = A[0]
A0tmpa += B[0]
A[0] = A0tmpa
A0tmpb = A[0]   # because C[1] = C[0] = 0
A0tmpb += B[1]
A[0] = A0tmpb

But suppose – in order to parallelize it – we wrote is as:
A0tmpa = A[0]
A0tmpb = A[0]   # because C[1] = C[0] = 0
A0tmpa += B[0]
A0tmpb += B[1]
A[0] = A0tmpa
A[0] = A0tmpb

In this case, you would get a different result, and probably not the
result you would expect or want.
Or just switching two lines would screw things up:
A0tmpa = A[0]
A0tmpa += B[0]
# A[0] = A0tmpa
# A0tmpb = A[0]   # because C[1] = C[0] = 0
A0tmpb = A[0]   # because C[1] = C[0] = 0
A[0] = A0tmpa
A0tmpb += B[1]
A[0] = A0tmpb

So the sequence of operations (which, if enforced, will break
parallelization) matters.
Here’s a script that shows – in a somewhat contrived way by
using floating-point round-off – that even if you perform each of
A[C[0]] += B[0] and A[C[1]] += B[1] “atomically,” the order
of the two atomic operations matters:
import torch
torch.__version__

delt = 16777216.0  # to get some round-off error

A = torch.FloatTensor ([[1, 2, 3], [4, 5, 6]])
D = A.clone()
B = torch.FloatTensor ([[delt, delt, delt], [-delt, -delt, -delt]])
C = torch.LongTensor ([0, 0])

A
A[C[0]] += B[0]
A
A[C[1]] += B[1]
A

D
D[C[1]] += B[1]
D
D[C[0]] += B[0]
D

And here is the output:
>>> import torch
>>> torch.__version__
'0.3.0b0+591e73e'
>>>
>>> delt = 16777216.0  # to get some round-off error
>>>
>>> A = torch.FloatTensor ([[1, 2, 3], [4, 5, 6]])
>>> D = A.clone()
>>> B = torch.FloatTensor ([[delt, delt, delt], [-delt, -delt, -delt]])
>>> C = torch.LongTensor ([0, 0])
>>>
>>> A

 1  2  3
 4  5  6
[torch.FloatTensor of size 2x3]

>>> A[C[0]] += B[0]
>>> A

 1.6777e+07  1.6777e+07  1.6777e+07
 4.0000e+00  5.0000e+00  6.0000e+00
[torch.FloatTensor of size 2x3]

>>> A[C[1]] += B[1]
>>> A

 0  2  4
 4  5  6
[torch.FloatTensor of size 2x3]

>>>
>>> D

 1  2  3
 4  5  6
[torch.FloatTensor of size 2x3]

>>> D[C[1]] += B[1]
>>> D

-1.6777e+07 -1.6777e+07 -1.6777e+07
 4.0000e+00  5.0000e+00  6.0000e+00
[torch.FloatTensor of size 2x3]

>>> D[C[0]] += B[0]
>>> D

 1  2  3
 4  5  6
[torch.FloatTensor of size 2x3]

And indeed, the results of iterating over the elements of C in
different orders differs.
One could imagine adding a function to pytorch that does what
you want so that you wouldn’t have to write an explicit loop.  But
because of the relatively strict sequencing required to give this
operation useful semantics, such a pytorch function would pretty
much have to be using a loop (or loop equivalent) under the hood,
so it’s unlikely doing so would buy you any significant speed-up.
Best.
K. Frank