zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-17 22:51:19 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix upon review

Browse Source
This commit is contained in:
Hyunsung Lee 2025-04-20 22:38:39 +09:00
parent f5dc7eaf6a
commit ed777f2ec6
3 changed files with 44 additions and 55 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
python/src/ops.cpp
View File

@ -5201,19 +5201,12 @@ void init_ops(nb::module_& m) {
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]));
}
if (nb::isinstance<mx::Shape>(shapes[i]) ||
nb::isinstance<nb::tuple>(shapes[i])) {
shape = nb::cast<mx::Shape>(shapes[i]);
} else {
throw std::invalid_argument(
"broadcast_shapes expects a sequence of shapes");
"broadcast_shapes expects a sequence of shapes (tuple or list of ints)");
}
shape_vec.push_back(shape);

44
python/tests/test_broadcast.py
View File

@ -1,44 +0,0 @@
# Copyright © 2025 Apple Inc.
import mlx.core
import mlx_tests
class TestBroadcast(mlx_tests.MLXTestCase):
def test_broadcast_shapes(self):
# Basic broadcasting
self.assertEqual(mlx.core.broadcast_shapes((1, 2, 3), (3,)), (1, 2, 3))
self.assertEqual(mlx.core.broadcast_shapes((4, 1, 6), (5, 6)), (4, 5, 6))
self.assertEqual(mlx.core.broadcast_shapes((5, 1, 4), (1, 3, 4)), (5, 3, 4))
# Multiple arguments
self.assertEqual(mlx.core.broadcast_shapes((1, 1), (1, 8), (7, 1)), (7, 8))
self.assertEqual(
mlx.core.broadcast_shapes((6, 1, 5), (1, 7, 1), (6, 7, 5)), (6, 7, 5)
)
# Same shapes
self.assertEqual(mlx.core.broadcast_shapes((3, 4, 5), (3, 4, 5)), (3, 4, 5))
# Single argument
self.assertEqual(mlx.core.broadcast_shapes((2, 3)), (2, 3))
# Empty shapes
self.assertEqual(mlx.core.broadcast_shapes((), ()), ())
self.assertEqual(mlx.core.broadcast_shapes((), (1,)), (1,))
self.assertEqual(mlx.core.broadcast_shapes((1,), ()), (1,))
# Broadcasting with zeroes
self.assertEqual(mlx.core.broadcast_shapes((0,), (0,)), (0,))
self.assertEqual(mlx.core.broadcast_shapes((1, 0, 5), (3, 1, 5)), (3, 0, 5))
self.assertEqual(mlx.core.broadcast_shapes((5, 0), (0, 5, 0)), (0, 5, 0))
# Error cases
with self.assertRaises(ValueError):
mlx.core.broadcast_shapes((3, 4), (4, 3))
with self.assertRaises(ValueError):
mlx.core.broadcast_shapes((2, 3, 4), (2, 5, 4))
with self.assertRaises(ValueError):
mlx.core.broadcast_shapes()

40
python/tests/test_ops.py
View File

@ -2935,5 +2935,45 @@ class TestOps(mlx_tests.MLXTestCase):
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()