[BUG] `mx.distributed.all_max` / `all_min` silently drop `NaN` (inconsistent with `mx.max`/`mx.min`)
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 28.5k
- Forks
- 2.3k
- Avg merge
- 3d 8h
- Merged PRs (30d)
- 62
Description
Describe the bug
The distributed cross-rank reduce uses std::max / std::min
(mlx/distributed/reduction_ops.h:17-35), which are not NaN-propagating
(std::max(nan, x) returns x when nan is the first argument). So if
exactly one rank has a NaN at a position, all_max / all_min return
the other ranks' non-NaN value and the NaN vanishes — the failure is
silently swallowed.
To Reproduce
import mlx.core as mx
# 2-rank group
nan = float("nan")
x = mx.array([1.0, nan if rank == 0 else 2.0, 3.0]) # rank 0 has NaN at idx 1
m = mx.distributed.all_max(x)
mx.eval(m) # -> array([1.0, 2.0, 3.0]) WRONG: should be [1.0, nan, 3.0]
Expected behavior
[rank 0] AMAX [1.0, 2.0, 3.0]
[rank 0] AMIN [1.0, 0.5, 3.0]
[rank 1] AMAX [1.0, 2.0, 3.0]
[rank 1] AMIN [1.0, 0.5, 3.0]
local mx.max([1,nan,3]) = nan (propagates NaN)
local mx.min([1,nan,3]) = nan (propagates NaN)
BUG REPRODUCED: all_max dropped the single-rank NaN (mx.max would propagate it).
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Read mlx/distributed/reduction_ops.h:17-35 and run the two-rank Python reproducer. Done when mx.distributed.all_max and all_min preserve a NaN at a position held by exactly one rank, matching mx.max and mx.min propagation on every rank.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- distributed-systems
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100