zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 09:21:16 +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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 10 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
mlx/ops.cpp
View File

@ -3175,6 +3175,10 @@ array scatter_axis(
throw std::invalid_argument(msg.str()); throw std::invalid_argument(msg.str());
} }
if (a.size() == 0) {
return a;
}
auto upd = astype(values, a.dtype(), s); auto upd = astype(values, a.dtype(), s);
// Squeeze leading singletons out of update // Squeeze leading singletons out of update

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) 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))) 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): def test_split(self):
a = mx.array([1, 2, 3]) a = mx.array([1, 2, 3])
splits = mx.split(a, 3) splits = mx.split(a, 3)