Is there an easy way to compute mean of every two tensors in a 2D tensor?

Suppose I have a tensor that looks like this:

a = torch.tensor([
[1. 2.],
[3. 4.],
[5. 6.],
[7. 8.]
])

Is there a way to compute the mean of every two row tensors (without overlap) in a without looping?

The expected output tensor would be:

torch.tensor([
[2. 3.], # mean of 1st and 2nd row tensors in `a`
[6. 7.]  # mean of 3rd and 4th row tensors in `a`
])

Something like this would work:

b = torch.stack([a_.mean(0) for a_ in a.split(2, 0)], 0)
print(b)
# tensor([[2., 3.],
#         [6., 7.]])

or using scatter_reduce_:

torch.zeros(2, 2).scatter_reduce_(0, torch.tensor([[0, 0], [0, 0], [1, 1], [1, 1]]), a, reduce="mean", include_self=False)

or index_put_ with accumulate=True and a scaling:

c = torch.zeros(2, 2)
c.index_put_((torch.tensor([0, 0, 1, 1]), ), a, accumulate=True)
c /= 2

Depending on your actual use case, how the indices are calculated, how large the tensor is etc. one approach might perform better than others.

1 Like