Skip to content

[BUG] mx.cummax / mx.cummin mishandle NaN: device-dependent (CPU ≠ GPU), leak ±∞, and disagree with NumPy #4044

Description

@JasonHonKL

Describe the bug
The Metal scan kernels for cumulative maximum / minimum
(mlx/backend/metal/kernels/scan.h, structs CumMax / CumMin) implement the
combine step with a plain ordered comparison that never propagates NaN:

To Reproduce

Include code snippet

import mlx.core as mx
import numpy as np

x = mx.array([1.0, 3.0, float("nan"), 5.0, 4.0])

print("np.maximum.accumulate:", np.maximum.accumulate(np.array(x)))
print("mx.cummax [gpu]      :", np.array(mx.cummax(x, stream=mx.gpu)))
print("mx.cummax [cpu]      :", np.array(mx.cummax(x, stream=mx.cpu)))

Expected behavior

np.maximum.accumulate: [ 1.  3. nan nan nan]
mx.cummax [gpu]      : [1. 3. 3. 5. 5.]      # NaN silently dropped
mx.cummax [cpu]      : [ 1.  3. nan  5.  5.]  # NaN kept at its slot, not propagated

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions