zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-19 00:04:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix sort (#1835)

Browse Source
This commit is contained in:
Awni Hannun
2025-02-05 17:16:27 -08:00
committed by GitHub
parent ca305afdbe
commit 9174606d4c
2 changed files with 23 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
python/tests/test_ops.py
View File

@@ -2010,6 +2010,15 @@ class TestOps(mlx_tests.MLXTestCase):
expected = mx.array([1, 3, 0, 2], dtype=mx.uint32)
self.assertTrue(mx.array_equal(out, expected))
# Test array with singleton dim
out = mx.sort(mx.array([1, 2, 3]), axis=0)
self.assertTrue(mx.array_equal(out, mx.array([1, 2, 3])))
x = np.random.uniform(size=(1, 4, 8, 1)).astype(np.float32)
y_np = np.sort(x, axis=-2)
y_mx = mx.sort(mx.array(x), axis=-2)
self.assertTrue(np.array_equal(y_np, y_mx))
def test_partition(self):
shape = (3, 4, 5)
for dtype in ("int32", "float32"):