[BUG] `mx.distributed.all_max` / `all_min` silently drop `NaN` (inconsistent with `mx.max`/`mx.min`)

[JasonHonKL]

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).

[erwinzhang7]

Reproduced on the JACCL backend over a four node mesh, and it's not only dropping
the NaN. Depending on which rank holds it, a real value from another rank goes
missing too.

Rank r contributes vals[r], one NaN among ordinary numbers:

NaN at	all_max	all_min	true max	true min
nowhere	4	1	4	1	correct
rank 0	4	2	4	2	correct
rank 1	4	3	4	1	rank 0's value is gone
rank 3	nan	nan	3	1	propagates instead of dropping

Three different behaviours from the same input depending on position: propagate,
drop cleanly, or drop the NaN and take another rank's value with it. The middle row
is the one that worries me. Ranks contributed 1, NaN, 3 and 4, and all_min returned
3. That is neither NaN nor the minimum, and nothing indicates anything went wrong.

That pattern is what a < b ? a : b produces: every comparison against NaN is false,
so which operand survives depends on the order the reduction tree combines them in,
and a real value can be discarded on the way.

The local path disagrees too: mx.max over the same four numbers
returns NaN, so distributed and local reductions do not currently agree about what a
NaN means, which is a second question on top of the corruption.

Repro is four ranks, one contributing float("nan") and the others rank + 1, then
compare dist.all_min against the minimum of the non-NaN values.

[JasonHonKL]

I can't really find the root cause of the issue. Solution provided is almost stack on top check which is neither elegant nor really solving the issue. Really curious where's the root cause. Have tried few times to figure it out but no clues.

[erwinzhang7]

@JasonHonKL the cause is that the reduce operator is not symmetric under NaN, which is
why checks at the call site do not fix it.

JACCL, mlx/distributed/jaccl/lib/jaccl/reduction_ops.h:

output[i] = (output[i] < input[i]) ? output[i] : input[i];   // MinOp

and the shared path in mlx/distributed/reduction_ops.h via std::min, which is
(b < a) ? b : a. Both return the other operand when the comparison is false, and
every comparison involving NaN is false:

• accumulator NaN, incoming 5: NaN < 5 false, take incoming, NaN dropped
• accumulator 5, incoming NaN: 5 < NaN also false, take incoming, NaN wins and the
  5 is gone

Two ranks, one NaN, both orderings, same two machines:

	NaN on rank 0	NaN on rank 1
jaccl	5.0	nan
ring	5.0	nan

Same op, same input, the answer decided by which rank holds the NaN. Both backends
behave the same, so it is one cause and not two. Guarding inputs cannot fix it: the op
is applied pairwise across ranks, and an operator that is not commutative under NaN
gives a result that depends on association order, which no caller controls.

The part that is not a NaN semantics question

On four ranks contributing 1, NaN, 3 and 4, all_min returns 3.

Not NaN, and not the minimum. The 1 was destroyed on the way past, by the second line
above, and the output carries no NaN to indicate anything happened. A job that reduces
a metric across ranks gets a plausible wrong number rather than a visibly poisoned one.

That is the difference between this and the rest of the family: #4274 and #4044 are
about which convention to follow, and this one silently corrupts data that was never
NaN. It reproduces on both backends and on any rank count above two, since it only
needs the NaN to be combined into an accumulator that already holds a real value.

On direction

If it is picked up, the semantics look settled rather than open: mx.max and
mx.cummax propagate today, #4044 fixed cummax/cummin that way, and #4274 argues
argmax/argmin should match. Propagating in the distributed ops would make them
agree with the local path, and it is order independent, which is the property actually
missing.

Why a fix sprawls, if that is useful for scoping

Only two of the four backends compute this themselves. ring uses std::min/std::max
and jaccl the ternary above, both in MLX. MPI hands off to MPI_MAX/MPI_MIN and
NCCL to ncclMax/ncclMin, so their behaviour is the library's. MPI could be given a
custom op through MPI_Op_create, but ncclRedOp_t is a fixed set with no general
custom op, so matching there means not using ncclAllReduce for min and max at all.

So a complete fix is three mechanisms and one backend that cannot express it, which
is probably how #4073 ended up as large as it did. A version scoped to ring and jaccl
is small and self contained, at the cost of those two agreeing with the local path
while MPI and NCCL still do whatever their library does.

Happy to test a patch on four M4 Pro minis across jaccl and ring.

[JasonHonKL]

Yea but any change will be massive.