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

Fix put_along_axis for empty arrays (#2181)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-05-13 14:27:53 -07:00
committed by GitHub
parent 8f3d208dce
commit 3aa9cf3f9e
2 changed files with 10 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/tests/test_ops.py
View File

@@ -1255,6 +1255,12 @@ class TestOps(mlx_tests.MLXTestCase):
np.put_along_axis(out_np, np.array(indices), np.array(update), axis=-2)
self.assertTrue(np.array_equal(out_np, np.array(out_mlx)))
a = mx.array([], mx.float32)
b = mx.put_along_axis(a, a, a, axis=None)
mx.eval(b)
self.assertEqual(b.size, 0)
self.assertEqual(b.shape, a.shape)
def test_split(self):
a = mx.array([1, 2, 3])
splits = mx.split(a, 3)