atomic_update in scatter

[zcbenz]

The scatter kernel uses atomic_update to set value, which I believe it is used to handle the case of duplicate indices:

mlx/mlx/backend/metal/kernels/scatter_axis.h

Line 51 in 659a519

op.atomic_update(out, upd[upd_idx], out_idx);

But considering duplicated indices in scatter is a non-deterministic behavior, shouldn't a normal assignment be enough?

And that is what PyTorch does:
https://github.com/pytorch/pytorch/blob/e6dccb036e185cd05c3859126bffd9644ba768b5/aten/src/ATen/native/cuda/ScatterGatherKernel.cu#L73-L81

And from PyTorch's documentation:

When indices are not unique, the behavior is non-deterministic (one of the values from src will be picked arbitrarily) and the gradient will be incorrect (it will be propagated to all locations in the source that correspond to the same index)!

[angeloskath]

That is partially true. We could replace the None reduction op with an assignment instead of an atomic store. However, scatter supports scatter add, scatter prod, scatter max etc and these need to be atomic operations to be correct.

[zcbenz]

Got it, thanks!