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 GPU sort for large arrays (#1285)

Browse Source 
* Fix GPU sort for large arrays
This commit is contained in:
Jagrit Digani
2024-07-24 14:37:10 -07:00
committed by GitHub
parent ebd7135b50
commit 7f914365fd
4 changed files with 34 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
python/tests/test_ops.py
View File

@@ -1840,6 +1840,15 @@ class TestOps(mlx_tests.MLXTestCase):
self.assertTrue(np.array_equal(c_np, c_mx))
self.assertEqual(b_mx.dtype, c_mx.dtype)
# Test very large array
if mx.default_device() == mx.gpu:
a_np = np.random.normal(20, 20, size=(2**22)).astype(np.float32)
a_mx = mx.array(a_np)
b_np = np.sort(a_np)
b_mx = mx.sort(a_mx)
self.assertTrue(np.array_equal(b_np, b_mx))
def test_partition(self):
shape = (3, 4, 5)
for dtype in ("int32", "float32"):