From bd200d6267a94ce7efe422517861f6c93ab483e5 Mon Sep 17 00:00:00 2001
From: Pranav Hari <39328505+pHequals7@users.noreply.github.com>
Date: Wed, 25 Mar 2026 07:16:15 +0530
Subject: [PATCH] Add output_shapes for AddMM (#3262)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>
---
 mlx/primitives.cpp           |  7 +++++++
 mlx/primitives.h             |  1 +
 python/tests/test_compile.py | 28 ++++++++++++++++++++++++++++
 3 files changed, 36 insertions(+)

diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp
index 8a70f894..566dcbf0 100644
--- a/mlx/primitives.cpp
+++ b/mlx/primitives.cpp
@@ -345,6 +345,13 @@ bool AddMM::is_equivalent(const Primitive& other) const {
   return (alpha_ == a_other.alpha_ && beta_ == a_other.beta_);
 }
 
+std::vector<Shape> AddMM::output_shapes(const std::vector<array>& inputs) {
+  // inputs are {a, b, c}, output shape is a's shape with last dim from b
+  auto out_shape = inputs[0].shape();
+  out_shape.back() = inputs[1].shape(-1);
+  return {std::move(out_shape)};
+}
+
 std::pair<std::vector<array>, std::vector<int>> AddMM::vmap(
     const std::vector<array>& inputs,
     const std::vector<int>& axes) {
diff --git a/mlx/primitives.h b/mlx/primitives.h
index a3ae63e5..ed580d6a 100644
--- a/mlx/primitives.h
+++ b/mlx/primitives.h
@@ -200,6 +200,7 @@ class AddMM : public UnaryPrimitive {
   DEFINE_NAME(AddMM)
 
   bool is_equivalent(const Primitive& other) const override;
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
   std::pair<float, float> state() const {
     return {alpha_, beta_};
   };
diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py
index 7db471cc..20f11452 100644
--- a/python/tests/test_compile.py
+++ b/python/tests/test_compile.py
@@ -884,6 +884,34 @@ class TestCompile(mlx_tests.MLXTestCase):
         fun = mx.compile(lambda a, b: a @ b, shapeless=True)
         self.assertTrue(mx.allclose(fun(a, b), a @ b))
 
+    def test_shapeless_compile_addmm(self):
+        def fun(c, a, b):
+            return mx.addmm(c, a, b)
+
+        cfun = mx.compile(fun, shapeless=True)
+
+        # First shape
+        c = mx.ones((2, 4))
+        a = mx.ones((2, 3))
+        b = mx.ones((3, 4))
+        self.assertTrue(mx.allclose(cfun(c, a, b), fun(c, a, b)))
+
+        # Different shape, same ranks — should not recompile
+        c = mx.ones((3, 5))
+        a = mx.ones((3, 6))
+        b = mx.ones((6, 5))
+        self.assertTrue(mx.allclose(cfun(c, a, b), fun(c, a, b)))
+
+        # With alpha and beta
+        fun2 = mx.compile(
+            lambda c, a, b: mx.addmm(c, a, b, alpha=2.0, beta=3.0), shapeless=True
+        )
+        c = mx.ones((2, 4))
+        a = mx.ones((2, 3))
+        b = mx.ones((3, 4))
+        expected = 3.0 * c + 2.0 * (a @ b)
+        self.assertTrue(mx.allclose(fun2(c, a, b), expected))
+
     def test_shapeless_compile_slice_update(self):
         def fun(x):
             x[2] = mx.array([3.0])