ml-explore / ml-explore/mlx

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

Open
#4,072 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug distributed low priority
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.