Possible bug in half-precision scatter

[noahstier]

I am getting different results for manually computing the mean of a half-precision tensor using torch.mean vs. torch_scatter.scatter.

Is this expected behavior?

import torch
import torch_scatter

src_float = torch.randn(1_000_000).float()
src_half = src_float.half()
idx = torch.zeros(len(src_float), dtype=torch.long)

result_float = torch_scatter.scatter(src_float, idx, reduce='mean')
result_half = torch_scatter.scatter(src_half, idx, reduce='mean')

result_float_manual = torch.mean(src_float)
result_half_manual = torch.mean(src_half)

print(result_float)
print(result_float_manual)
print(result_half)
print(result_half_manual)

prints:

tensor([-0.0014])
tensor(-0.0014)
tensor([-0.7241], dtype=torch.float16)
tensor(-0., dtype=torch.float16)

[rusty1s]

Are you sure result_half_manual provides the correct result? To me, this looks more like overfloating problem since 1_000_000 (the number of elements) is not representable by torch.half. For example,

src_half = src_float.half().abs()
result_half_manual = torch.mean(src_half)

produces tensor(nan, dtype=torch.float16). Everything looks good as well in case you reduce the number of elements to, e.g., 1_000.

[noahstier]

Using 10_000 elements is a better example:

tensor([0.0024])
tensor(0.0024)
tensor([0.0119], dtype=torch.float16)
tensor(0.0024, dtype=torch.float16)

Here, result_half_manual is correct but result_half is not. Even in the case of 1_000_000 elements though, the fact that the results are different is what drew my eye, even if neither is correct in that case.

[rusty1s]

Indeed. Seems to be a bug in torch.scatter_add, where

N = 10000
src = torch.ones(N, dtype=torch.half)
out = torch.zeros(1, dtype=torch.half)
index = torch.zeros(N, dtype=torch.long)
out.scatter_add_(0, index, src)

cannot produce an output larger than 2048.

[github-actions]

This issue had no activity for 6 months. It will be closed in 2 weeks unless there is some new activity. Is this issue already resolved?