From 1b1c56352adf5e505bf676c4366e5db474e1ba00 Mon Sep 17 00:00:00 2001
From: Sheldon Aristide <44308312+Aristide021@users.noreply.github.com>
Date: Tue, 24 Mar 2026 20:02:31 -0400
Subject: [PATCH] Fix moved-from shape bug in broadcast_arrays causing vmap bus
 error (#3310)

---
 mlx/ops.cpp               |  2 +-
 python/tests/test_vmap.py |  6 ++++++
 tests/vmap_tests.cpp      | 12 ++++++++++++
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/mlx/ops.cpp b/mlx/ops.cpp
index 935bbd6e..ef792cd6 100644
--- a/mlx/ops.cpp
+++ b/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}));
diff --git a/python/tests/test_vmap.py b/python/tests/test_vmap.py
index 30c85de0..5d60c0e1 100644
--- a/python/tests/test_vmap.py
+++ b/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)
diff --git a/tests/vmap_tests.cpp b/tests/vmap_tests.cpp
index 36198284..c83a10e7 100644
--- a/tests/vmap_tests.cpp
+++ b/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,