[Feature Request] XLA (and NumPy) compatible sorting for NaNs

[polvalente]

I'm starting this as a discussion because I'm not sure if this has been considered and there wasn't an issue template for feature requests.

In [1]: import numpy as np
In [5]: np.sort(t)
Out[5]: array([ 0.,  1.,  2., nan])
In [6]: import mlx.core as mx
In [7]: t = mx.array([1.0, np.nan, 2.0, 0.0])
In [8]: mx.sort(t)
Out[8]: array([1, nan, 0, 2], dtype=float32)

In short, NaNs in numpy will be sorted to the end while MLX seems to not care about them. I can see the reasoning behind this, but it would be nice if we could have an option to control this.

I believe this would affect max, min, argmax, argmin, argsort and sort. Not sure if there are any other similar functions.

[awni]

Well aside whether or not we can provide that option.. it looks like the NaN breaks the sort which is perhaps unexpected.

[awni]

Every other framework I tested puts NaNs to the end. I think it's reasonable to make that change in MLX and probably not too difficult to implement.

[polvalente]

Yeah, NaNs are always special-cased when I need to implement things with comparisons in Nx. Kinda annoying to be honest, specially given the "always returns false for any comparison" behavior.