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

Fix Split::vmap (#1845)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-02-08 09:22:13 -08:00
committed by GitHub
parent 1c0c118f7c
commit 9eb7d7362f
2 changed files with 15 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
mlx/primitives.cpp
View File

@@ -4273,7 +4273,9 @@ std::pair<std::vector<array>, std::vector<int>> Split::vmap(
const std::vector<array>& inputs,
const std::vector<int>& axes) {
int axis_left = axes[0] >= 0 && axes[0] <= axis_;
return {{split(inputs[0], indices_, axis_ + axis_left, stream())}, axes};
auto output = split(inputs[0], indices_, axis_ + axis_left, stream());
std::vector<int> output_axes(output.size(), axes[0]);
return {std::move(output), std::move(output_axes)};
}
std::vector<array> Split::vjp(

12
python/tests/test_vmap.py
View File

@@ -596,6 +596,18 @@ class TestVmap(mlx_tests.MLXTestCase):
out = mx.vmap(fun, in_axes=(None, 1, 1))(a, idx, upd)
self.assertEqual(out.shape, (4, 5, 1))
def test_vmap_split_vmap(self):
def fun(x):
a, b = mx.split(x, 2, 1)
return mx.concatenate([b, a], 1)
x = mx.ones((5, 6, 7))
y = mx.ones((5, 4, 6, 7))
fx = fun(x)
fy = mx.vmap(fun, in_axes=1)(y)
self.assertEqual(fx.shape, (5, 6, 7))
self.assertEqual(fy.shape, (4, 5, 6, 7))
if __name__ == "__main__":
unittest.main()