L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix sort NaN handling for float16 and bfloat16 (#3269)

Browse Source
This commit is contained in:
Long Yixing
2026-03-20 06:19:41 +08:00
committed by GitHub
parent 5fa1a8d59f
commit 82809ebd12
4 changed files with 26 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tests/test_ops.py
+15 -3
View File
@@ -3311,11 +3311,23 @@ class TestOps(mlx_tests.MLXTestCase):
mx.broadcast_shapes()
def test_sort_nan(self):
x = mx.array([3.0, mx.nan, 2.0, 0.0])
expected = mx.array([0.0, 2.0, 3.0, mx.nan])
self.assertTrue(mx.array_equal(mx.sort(x), expected, equal_nan=True))
for dtype in [mx.float32, mx.float16, mx.bfloat16]:
with self.subTest(dtype=dtype):
x = mx.array([3.0, mx.nan, 2.0, 0.0], dtype=dtype)
expected = mx.array([0.0, 2.0, 3.0, mx.nan], dtype=dtype)
self.assertTrue(mx.array_equal(mx.sort(x), expected, equal_nan=True))
x = mx.array([3.0, mx.nan, 2.0, 0.0]) + 1j * mx.array([1.0] * 4)
def test_argsort_nan(self):
for dtype in [mx.float32, mx.float16, mx.bfloat16]:
with self.subTest(dtype=dtype):
x = mx.array([3.0, mx.nan, 2.0, 0.0], dtype=dtype)
expected = mx.array([0.0, 2.0, 3.0, mx.nan], dtype=dtype)
indices = mx.argsort(x)
sorted_x = mx.take(x, indices)
self.assertTrue(mx.array_equal(sorted_x, expected, equal_nan=True))
def test_to_from_fp8(self):
vals = mx.array(
[448, 256, 192, 128, 96, 64, 48, 32, 24, 16, 12, 8, 6, 4, 3, 2, 0.015625]