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

Fix multi-block sort stride management (#1169)

Browse Source 
* Fix multi-block sort stride management

* Add seed to tests
This commit is contained in:
Jagrit Digani
2024-05-31 11:10:54 -07:00
committed by GitHub
parent 9f0df51f8d
commit 76b6cece46
2 changed files with 33 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
python/tests/test_ops.py
View File

@@ -1754,6 +1754,9 @@ class TestOps(mlx_tests.MLXTestCase):
self.assertTrue(np.array_equal(d_np, d_mx))
self.assertEqual(c_mx.dtype, mx.uint32)
# Set random seed
np.random.seed(0)
# Test multi-block sort
a_np = np.random.normal(size=(32769,)).astype(np.float32)
a_mx = mx.array(a_np)
@@ -1764,6 +1767,25 @@ class TestOps(mlx_tests.MLXTestCase):
self.assertTrue(np.array_equal(b_np, b_mx))
self.assertEqual(b_mx.dtype, a_mx.dtype)
# Test multi-dum multi-block sort
a_np = np.random.normal(size=(2, 4, 32769)).astype(np.float32)
a_mx = mx.array(a_np)
b_np = np.sort(a_np, axis=-1)
b_mx = mx.sort(a_mx, axis=-1)
self.assertTrue(np.array_equal(b_np, b_mx))
self.assertEqual(b_mx.dtype, a_mx.dtype)
a_np = np.random.normal(size=(2, 32769, 4)).astype(np.float32)
a_mx = mx.array(a_np)
b_np = np.sort(a_np, axis=1)
b_mx = mx.sort(a_mx, axis=1)
self.assertTrue(np.array_equal(b_np, b_mx))
self.assertEqual(b_mx.dtype, a_mx.dtype)
def test_partition(self):
shape = (3, 4, 5)
for dtype in ("int32", "float32"):