Given a 3D tensor A of size

(n, i, j)

And an index tensor of size n with a pattern like:

[1, 2, 2, 3, 3, 3, 4, 4, 4, 4 …]

I want to sum over the n dimensions of A according to the index tensor.

For example, if A is (6, 2, 2)

```
tensor([[[0.4, 0.4],
[0.6, 0.3]],
[[0.9, 0.3],
[0.5, 0.6]],
[[0.4, 0.1],
[0.4, 0.8]],
[[0.6, 0.3],
[0.7, 0.1]],
[[0.0, 0.9],
[0.4, 0.7]],
[[0.0, 0.4],
[0.5, 0.7]]])
```

I want to sum A[1] + A[2] and A[3] + A[4] + A[5].

The result has size (3,2,2) and is equal to:

```
tensor([[[0.4, 0.4],
[0.6, 0.3]],
[[1.3, 0.4],
[0.9, 1.4]],
[[0.6, 1.6],
[1.6, 1.5]]])
```

Which is the best way to do this, for large n ?