Fix shapeless matmul with dynamic batch dimensions

mlx/ops.cpp

@@ -3352,9 +3352,14 @@ array matmul(
   }
 
   // We can batch the multiplication by reshaping a
-  if (in_a.ndim() > 2 && in_b.ndim() <= 2) {
+  bool flattened_a = false;
+
+  // Avoid flatten/unflatten during dynamic tracing because unflatten stores the
+  // trace-time shape, which breaks shapeless replay with dynamic batch sizes.
+  if (in_a.ndim() > 2 && in_b.ndim() <= 2 && !detail::in_dynamic_tracing()) {
     a = flatten(a, 0, -2, s);
-  } else if (in_b.ndim() > 2) {
+    flattened_a = true;
+  } else if (in_a.ndim() > 2 || in_b.ndim() > 2) {
     std::tie(a, b) = broadcast_arrays(a, b, {-2, -1}, s);
   }
 
@@ -3366,7 +3371,8 @@ array matmul(
       out_type,
       std::make_shared<Matmul>(to_stream(s)),
       {a, b});
-  if (in_a.ndim() > 2 && in_b.ndim() <= 2) {
+
+  if (flattened_a) {
     auto orig_shape = in_a.shape();
     orig_shape.pop_back();
     out = unflatten(out, 0, std::move(orig_shape), s);

python/tests/test_export_import.py

@@ -704,6 +704,50 @@ def fun(x, y, z):
         imported = mx.import_function(path)
         self.assertTrue(mx.array_equal(imported(x, y, z)[0], fun(x, y, z)))
 
+    def test_export_matmul_shapeless_mid_dim(self):
+        path = os.path.join(self.test_dir, "matmul_shapeless.mlxfn")
+
+        E, H = 64, 17
+        arr = mx.arange(E * H, dtype=mx.float32).reshape((E, H)) * (1.0 / (E * H))
+
+        def fn(x):
+            return mx.matmul(x, arr)
+
+        sample = mx.zeros((1, 40, E), dtype=mx.float32)
+        mx.export_function(path, fn, sample, shapeless=True)
+        imported = mx.import_function(path)
+
+        for seq_len in (40, 248, 623):
+            with self.subTest(seq_len=seq_len):
+                x = mx.arange(seq_len * E, dtype=mx.float32).reshape((1, seq_len, E))
+                expected = fn(x)
+                (y,) = imported(x)
+                self.assertEqual(y.shape, (1, seq_len, H))
+                self.assertTrue(mx.allclose(y, expected))
+
+    def test_export_matmul_shapeless_batch_and_mid_dim(self):
+        path = os.path.join(self.test_dir, "matmul_shapeless_batch.mlxfn")
+
+        B, E, H = 2, 32, 8
+        arr = mx.arange(E * H, dtype=mx.float32).reshape((E, H)) * (1.0 / (E * H))
+
+        def fn(x):
+            return mx.matmul(x, arr)
+
+        sample = mx.zeros((B, 10, E), dtype=mx.float32)
+        mx.export_function(path, fn, sample, shapeless=True)
+        imported = mx.import_function(path)
+
+        for seq_len in (10, 50, 100):
+            with self.subTest(seq_len=seq_len):
+                x = mx.arange(B * seq_len * E, dtype=mx.float32).reshape(
+                    (B, seq_len, E)
+                )
+                expected = fn(x)
+                (y,) = imported(x)
+                self.assertEqual(y.shape, (B, seq_len, H))
+                self.assertTrue(mx.allclose(y, expected))
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()