Different outputs of mlx.core.max using cpu/gpu

[Redempt1onzzZZ]

Describe the bug Different(maybe wrong) outputs of mlx.core.max using cpu/gpu compared to numpy

To Reproduce

import mlx.core as mx
import numpy as np
array_mlx = mx.array([[1, 2, 3], [float('nan'), 5, 6]], dtype=mx.float32)

mx.set_default_device(mx.gpu)
max_gpu = mx.max(array_mlx, axis=0)
print("Max on GPU:", max_gpu)

mx.set_default_device(mx.cpu)
max_cpu = mx.max(array_mlx, axis=0)
print("Max on CPU:", max_cpu)

a_np = np.array([[1, 2, 3], [float('nan'), 5, 6]], dtype=np.float32)
max_np = np.max(a_np, axis=0)
print("Max using NumPy:", max_np)

Desktop (please complete the following information):

• OS Version: MacOS 14.2.1
• Version 0.0.10

[awni]

Actually this one is not so easy to fix as we use the simd instructions in Metal. The NaN are being suppressed there.