Fix moved-from shape bug in broadcast_arrays causing vmap bus error (#3310)

mlx/ops.cpp

@@ -1659,7 +1659,7 @@ std::vector<array> broadcast_arrays(
         outputs.push_back(in);
       } else {
         outputs.push_back(array(
-            std::move(out_shape),
+            out_shape,
             in.dtype(),
             std::make_shared<Broadcast>(to_stream(s), out_shape),
             {in}));

python/tests/test_vmap.py

@@ -595,6 +595,12 @@ class TestVmap(mlx_tests.MLXTestCase):
         out = mx.vmap(fun, in_axes=(None, 0))(a, idx)
         self.assertEqual(out.shape, (4, 2, 1))
 
+        a = mx.zeros((4, 5, 3))
+        idx = mx.zeros((2, 2, 1, 3), mx.int32)
+
+        out = mx.vmap(fun, in_axes=(None, 0))(a, idx)
+        self.assertEqual(out.shape, (2, 2, 5, 3))
+
     def test_vmap_put_along_axis(self):
         a = mx.zeros((4, 5, 1))
         idx = mx.ones((2, 4, 1), mx.int32)

tests/vmap_tests.cpp

@@ -392,6 +392,18 @@ TEST_CASE("test vmap gather") {
   }
 }
 
+TEST_CASE("test vmap take_along_axis with unmapped input and mapped index") {
+  auto fun = [](std::vector<array> inputs) {
+    return std::vector<array>{take_along_axis(inputs[0], inputs[1], 0)};
+  };
+
+  auto a = reshape(arange(60), {4, 5, 3});
+  auto idx = zeros({2, 2, 1, 3}, int32);
+
+  auto out = vmap(fun, {-1, 0})({a, idx})[0];
+  CHECK_EQ(out.shape(), Shape{2, 2, 5, 3});
+}
+
 TEST_CASE("test vmap scatter") {
   auto make_scatter_fn = [](const std::vector<array>& indices,
                             const array& updates,