Repository: ml-explore/mlx
Pull request #2091: Add broadcast_shapes in python API
Author: login: ita9naiwa
Created: 2025-04-19T06:43:17Z | Merged: 2025-04-23T01:57:39Z

## Proposed changes

Add a `broadcast_shapes` function to the Python API. (related #2089)

Since this operation is simple (working only on a list of shape tuples), a pure Python implementation is preferable to reduce C++ binding boilerplate and keep the API surface minimal.

If necessary, I’m open to switching to a C++-exposed version.

## Checklist

Put an `x` in the boxes that apply.

- [x] I have read the [CONTRIBUTING](https://github.com/ml-explore/mlx/blob/main/CONTRIBUTING.md) document
- [x] I have run `pre-commit run --all-files` to format my code / installed pre-commit prior to committing changes
- [x] I have added tests that prove my fix is effective or that my feature works
- [x] I have updated the necessary documentation (if needed)



Diff:
diff --git a/python/src/ops.cpp b/python/src/ops.cpp
index f98aa80aac..5969c50524 100644
--- a/python/src/ops.cpp
+++ b/python/src/ops.cpp
@@ -5189,4 +5189,46 @@ void init_ops(nb::module_& m) {
       Returns:
         array: The row or col contiguous output.
     )pbdoc");
+  m.def(
+      "broadcast_shapes",
+      [](const nb::args& shapes) {
+        if (shapes.size() == 0)
+          throw std::invalid_argument(
+              "[broadcast_shapes] Must provide at least one shape.");
+
+        mx::Shape result = nb::cast<mx::Shape>(shapes[0]);
+        for (size_t i = 1; i < shapes.size(); ++i) {
+          if (!nb::isinstance<mx::Shape>(shapes[i]) &&
+              !nb::isinstance<nb::tuple>(shapes[i]))
+            throw std::invalid_argument(
+                "[broadcast_shapes] Expects a sequence of shapes (tuple or list of ints).");
+          result = mx::broadcast_shapes(result, nb::cast<mx::Shape>(shapes[i]));
+        }
+
+        return nb::tuple(nb::cast(result));
+      },
+      nb::sig("def broadcast_shapes(*shapes: Sequence[int]) -> Tuple[int]"),
+      R"pbdoc(
+        Broadcast shapes.
+
+        Returns the shape that results from broadcasting the supplied array shapes
+        against each other.
+
+        Args:
+            *shapes (Sequence[int]): The shapes to broadcast.
+
+        Returns:
+            tuple: The broadcasted shape.
+
+        Raises:
+            ValueError: If the shapes cannot be broadcast.
+
+        Example:
+            >>> mx.broadcast_shapes((1,), (3, 1))
+            (3, 1)
+            >>> mx.broadcast_shapes((6, 7), (5, 6, 1), (7,))
+            (5, 6, 7)
+            >>> mx.broadcast_shapes((5, 1, 4), (1, 3, 1))
+            (5, 3, 4)
+      )pbdoc");
 }
diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py
index 4fcb31f18a..5008abb3f5 100644
--- a/python/tests/test_ops.py
+++ b/python/tests/test_ops.py
@@ -2935,5 +2935,45 @@ def test_slice_with_negative_stride(self):
         self.assertTrue(mx.array_equal(out[-1, :], a[0, :]))
 
 
+class TestBroadcast(mlx_tests.MLXTestCase):
+    def test_broadcast_shapes(self):
+        # Basic broadcasting
+        self.assertEqual(mx.broadcast_shapes((1, 2, 3), (3,)), (1, 2, 3))
+        self.assertEqual(mx.broadcast_shapes((4, 1, 6), (5, 6)), (4, 5, 6))
+        self.assertEqual(mx.broadcast_shapes((5, 1, 4), (1, 3, 4)), (5, 3, 4))
+
+        # Multiple arguments
+        self.assertEqual(mx.broadcast_shapes((1, 1), (1, 8), (7, 1)), (7, 8))
+        self.assertEqual(
+            mx.broadcast_shapes((6, 1, 5), (1, 7, 1), (6, 7, 5)), (6, 7, 5)
+        )
+
+        # Same shapes
+        self.assertEqual(mx.broadcast_shapes((3, 4, 5), (3, 4, 5)), (3, 4, 5))
+
+        # Single argument
+        self.assertEqual(mx.broadcast_shapes((2, 3)), (2, 3))
+
+        # Empty shapes
+        self.assertEqual(mx.broadcast_shapes((), ()), ())
+        self.assertEqual(mx.broadcast_shapes((), (1,)), (1,))
+        self.assertEqual(mx.broadcast_shapes((1,), ()), (1,))
+
+        # Broadcasting with zeroes
+        self.assertEqual(mx.broadcast_shapes((0,), (0,)), (0,))
+        self.assertEqual(mx.broadcast_shapes((1, 0, 5), (3, 1, 5)), (3, 0, 5))
+        self.assertEqual(mx.broadcast_shapes((5, 0), (0, 5, 0)), (0, 5, 0))
+
+        # Error cases
+        with self.assertRaises(ValueError):
+            mx.broadcast_shapes((3, 4), (4, 3))
+
+        with self.assertRaises(ValueError):
+            mx.broadcast_shapes((2, 3, 4), (2, 5, 4))
+
+        with self.assertRaises(ValueError):
+            mx.broadcast_shapes()
+
+
 if __name__ == "__main__":
     unittest.main()


Reviews:
user: login: awni
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-20T13:26:03Z
user: login: awni
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-20T13:27:05Z
user: login: ita9naiwa
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-20T13:40:01Z
user: login: ita9naiwa
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-20T13:40:41Z
user: login: ita9naiwa
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-20T13:40:47Z
user: login: awni
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-21T13:05:55Z
user: login: awni
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-21T13:06:14Z
user: login: awni
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-21T13:07:12Z
user: login: awni
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-21T13:07:33Z
user: login: awni
body: Very nice, thanks!
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-21T20:55:24Z

Review comments:
diff_hunk: @@ -5189,4 +5189,65 @@ void init_ops(nb::module_& m) {
       Returns:
         array: The row or col contiguous output.
     )pbdoc");
+  m.def(
+      "broadcast_shapes",
+      [](const nb::args& shapes) {
+        if (shapes.size() == 0)
+          throw std::invalid_argument(
+              "broadcast_shapes expects a sequence of shapes");
+
+        std::vector<mx::Shape> shape_vec;
+        shape_vec.reserve(shapes.size());
+
+        for (size_t i = 0; i < shapes.size(); ++i) {
+          mx::Shape shape;
+          if (nb::isinstance<nb::tuple>(shapes[i])) {
+            nb::tuple t = nb::cast<nb::tuple>(shapes[i]);
+            for (size_t j = 0; j < t.size(); ++j) {
+              shape.push_back(nb::cast<int>(t[j]));
+            }
+          } else if (nb::isinstance<nb::list>(shapes[i])) {
+            nb::list l = nb::cast<nb::list>(shapes[i]);
+            for (size_t j = 0; j < l.size(); ++j) {
+              shape.push_back(nb::cast<int>(l[j]));
+            }
path: python/src/ops.cpp
user: login: awni
body: These loops are not necessary. You should be able to get a `mx::Shape` from the `list` or `tuple` in a single line.
created_at: 2025-04-20T13:26:03Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,44 @@
+# Copyright © 2025 Apple Inc.
+
+import mlx.core
+import mlx_tests
+
+
+class TestBroadcast(mlx_tests.MLXTestCase):
+    def test_broadcast_shapes(self):
path: python/tests/test_broadcast.py
user: login: awni
body: I don't think we need a separate file for this test since I don't see us adding anything here. Just put it in the ops test for now.
created_at: 2025-04-20T13:27:05Z
author_association: MEMBER
diff_hunk: @@ -5189,4 +5189,65 @@ void init_ops(nb::module_& m) {
       Returns:
         array: The row or col contiguous output.
     )pbdoc");
+  m.def(
+      "broadcast_shapes",
+      [](const nb::args& shapes) {
+        if (shapes.size() == 0)
+          throw std::invalid_argument(
+              "broadcast_shapes expects a sequence of shapes");
+
+        std::vector<mx::Shape> shape_vec;
+        shape_vec.reserve(shapes.size());
+
+        for (size_t i = 0; i < shapes.size(); ++i) {
+          mx::Shape shape;
+          if (nb::isinstance<nb::tuple>(shapes[i])) {
+            nb::tuple t = nb::cast<nb::tuple>(shapes[i]);
+            for (size_t j = 0; j < t.size(); ++j) {
+              shape.push_back(nb::cast<int>(t[j]));
+            }
+          } else if (nb::isinstance<nb::list>(shapes[i])) {
+            nb::list l = nb::cast<nb::list>(shapes[i]);
+            for (size_t j = 0; j < l.size(); ++j) {
+              shape.push_back(nb::cast<int>(l[j]));
+            }
path: python/src/ops.cpp
user: login: ita9naiwa
body: Thanks!

created_at: 2025-04-20T13:40:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,44 @@
+# Copyright © 2025 Apple Inc.
+
+import mlx.core
+import mlx_tests
+
+
+class TestBroadcast(mlx_tests.MLXTestCase):
+    def test_broadcast_shapes(self):
path: python/tests/test_broadcast.py
user: login: ita9naiwa
body: Right. I moved `broadcast_shapes` test into `test_ops.py`
created_at: 2025-04-20T13:40:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -5189,4 +5189,65 @@ void init_ops(nb::module_& m) {
       Returns:
         array: The row or col contiguous output.
     )pbdoc");
+  m.def(
+      "broadcast_shapes",
+      [](const nb::args& shapes) {
+        if (shapes.size() == 0)
+          throw std::invalid_argument(
+              "broadcast_shapes expects a sequence of shapes");
+
+        std::vector<mx::Shape> shape_vec;
+        shape_vec.reserve(shapes.size());
+
+        for (size_t i = 0; i < shapes.size(); ++i) {
+          mx::Shape shape;
+          if (nb::isinstance<nb::tuple>(shapes[i])) {
+            nb::tuple t = nb::cast<nb::tuple>(shapes[i]);
+            for (size_t j = 0; j < t.size(); ++j) {
+              shape.push_back(nb::cast<int>(t[j]));
+            }
+          } else if (nb::isinstance<nb::list>(shapes[i])) {
+            nb::list l = nb::cast<nb::list>(shapes[i]);
+            for (size_t j = 0; j < l.size(); ++j) {
+              shape.push_back(nb::cast<int>(l[j]));
+            }
path: python/src/ops.cpp
user: login: ita9naiwa
body: fixed!

created_at: 2025-04-20T13:40:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -5189,4 +5189,54 @@ void init_ops(nb::module_& m) {
       Returns:
         array: The row or col contiguous output.
     )pbdoc");
+  m.def(
+      "broadcast_shapes",
+      [](const nb::args& shapes) {
+        if (shapes.size() == 0)
+          throw std::invalid_argument(
+              "broadcast_shapes expects a sequence of shapes");
path: python/src/ops.cpp
user: login: awni
body: Could you make that error message a little more clear?
```suggestion
              "[broadcast_shapes] Must provide at least one shape.");
```
created_at: 2025-04-21T13:05:55Z
author_association: MEMBER
diff_hunk: @@ -5189,4 +5189,54 @@ void init_ops(nb::module_& m) {
       Returns:
         array: The row or col contiguous output.
     )pbdoc");
+  m.def(
+      "broadcast_shapes",
+      [](const nb::args& shapes) {
+        if (shapes.size() == 0)
+          throw std::invalid_argument(
+              "broadcast_shapes expects a sequence of shapes");
+
+        std::vector<mx::Shape> shape_vec;
+        shape_vec.reserve(shapes.size());
+
+        for (size_t i = 0; i < shapes.size(); ++i) {
+          if (!nb::isinstance<mx::Shape>(shapes[i]) &&
+              !nb::isinstance<nb::tuple>(shapes[i]))
+            throw std::invalid_argument(
+                "broadcast_shapes expects a sequence of shapes (tuple or list of ints)");
path: python/src/ops.cpp
user: login: awni
body: ```suggestion
                "[broadcast_shapes] Expects a sequence of shapes (tuple or list of ints).");
```
created_at: 2025-04-21T13:06:13Z
author_association: MEMBER
diff_hunk: @@ -5189,4 +5189,54 @@ void init_ops(nb::module_& m) {
       Returns:
         array: The row or col contiguous output.
     )pbdoc");
+  m.def(
+      "broadcast_shapes",
+      [](const nb::args& shapes) {
+        if (shapes.size() == 0)
+          throw std::invalid_argument(
+              "broadcast_shapes expects a sequence of shapes");
+
+        std::vector<mx::Shape> shape_vec;
+        shape_vec.reserve(shapes.size());
+
+        for (size_t i = 0; i < shapes.size(); ++i) {
+          if (!nb::isinstance<mx::Shape>(shapes[i]) &&
+              !nb::isinstance<nb::tuple>(shapes[i]))
+            throw std::invalid_argument(
+                "broadcast_shapes expects a sequence of shapes (tuple or list of ints)");
+
+          shape_vec.push_back(nb::cast<mx::Shape>(shapes[i]));
path: python/src/ops.cpp
user: login: awni
body: What if we just broadcast here instead of making a temporary vector of shapes?

```suggestion
          auto shape = nb::cast<mx::Shape>(shapes[i]);
          result = i == 0 ? shape : mx::broadcast_shapes(result, shape);
```
created_at: 2025-04-21T13:07:12Z
author_association: MEMBER
diff_hunk: @@ -5189,4 +5189,54 @@ void init_ops(nb::module_& m) {
       Returns:
         array: The row or col contiguous output.
     )pbdoc");
+  m.def(
+      "broadcast_shapes",
+      [](const nb::args& shapes) {
+        if (shapes.size() == 0)
+          throw std::invalid_argument(
+              "broadcast_shapes expects a sequence of shapes");
+
+        std::vector<mx::Shape> shape_vec;
+        shape_vec.reserve(shapes.size());
+
+        for (size_t i = 0; i < shapes.size(); ++i) {
+          if (!nb::isinstance<mx::Shape>(shapes[i]) &&
+              !nb::isinstance<nb::tuple>(shapes[i]))
+            throw std::invalid_argument(
+                "broadcast_shapes expects a sequence of shapes (tuple or list of ints)");
+
+          shape_vec.push_back(nb::cast<mx::Shape>(shapes[i]));
+        }
+
+        mx::Shape result = shape_vec[0];
+        for (size_t i = 1; i < shape_vec.size(); ++i)
+          result = mx::broadcast_shapes(result, shape_vec[i]);
path: python/src/ops.cpp
user: login: awni
body: Then remove this and put the initializer further up.
created_at: 2025-04-21T13:07:32Z
author_association: MEMBER

PR comments:
user: login: awni
created_at: 2025-04-19T13:09:46Z
body: > If necessary, I’m open to switching to a C++-exposed version.

Yes please switch to that. This should be a function in `mlx.core` and there's no need to duplicate the broadcasting logic in Python.
author_association: MEMBER
user: login: ita9naiwa
created_at: 2025-04-19T21:56:13Z
body: Okay.
I'll make following changes soon.
- place `broadcast_shapes` in `mlx.core`
- use C++ api
thanks for your comment!
author_association: CONTRIBUTOR
user: login: ita9naiwa
created_at: 2025-04-20T03:59:33Z
body: fixed! 
@awni could you please review this?
author_association: CONTRIBUTOR
user: login: ita9naiwa
created_at: 2025-04-20T13:40:59Z
body: fixed! Thanks for quick review!
author_association: CONTRIBUTOR
user: login: ita9naiwa
created_at: 2025-04-21T20:54:09Z
body: Done! Thanks, now code looks much clear!
author_association: CONTRIBUTOR
user: login: ita9naiwa
created_at: 2025-04-21T20:56:59Z
body: @awni Sorry I made a last minute change... can you check this??

I think my previous commit didn't look safe to me? (although test passing)

prev one:
```c++
mx::Shape result;
for (size_t i = 0; i < shapes.size(); ++i) {
  if (!nb::isinstance<mx::Shape>(shapes[i]) &&
      !nb::isinstance<nb::tuple>(shapes[i]))
    throw std::invalid_argument(
        "[broadcast_shapes] Expects a sequence of shapes (tuple or list of ints).");
  result = mx::broadcast_shapes(result, nb::cast<mx::Shape>(shapes[i]));
}
```

current:
```c++
mx::Shape result = nb::cast<mx::Shape>(shapes[0]);
for (size_t i = 1; i < shapes.size(); ++i) {
  if (!nb::isinstance<mx::Shape>(shapes[i]) &&
      !nb::isinstance<nb::tuple>(shapes[i]))
    throw std::invalid_argument(
        "[broadcast_shapes] Expects a sequence of shapes (tuple or list of ints).");
  result = mx::broadcast_shapes(result, nb::cast<mx::Shape>(shapes[i]));
}
```
author_association: CONTRIBUTOR
user: login: awni
created_at: 2025-04-22T23:19:45Z
body: The tests are running. Once they clear we can merge this.
author_association: MEMBER