KFrank:

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.


… And this operation is called scatter_add, it uses atomicAdd on cuda, on cpu it is a loop, but c++ loop != python loop.
A.scatter_add(0, C.view(2,-1).expand(-1,3), B)

PS about your rounding example, many cuda ops are non-deterministic, so we live with it, trying not to sum billions with one in float32.