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
Describe the bug
The Metal scan kernels for cumulative maximum / minimum
(
mlx/backend/metal/kernels/scan.h, structsCumMax/CumMin) implement thecombine step with a plain ordered comparison that never propagates
NaN:To Reproduce
Include code snippet
Expected behavior