mirror of
https://github.com/ml-explore/mlx.git
synced 2025-06-25 09:51:17 +08:00
45 lines
1.6 KiB
Python
45 lines
1.6 KiB
Python
|
# Copyright © 2023 Apple Inc.
|
||
|
|
|
||
|
|
import mlx.utils
|
||
|
|
import mlx_tests
|
||
|
|
|
||
|
|
|
||
|
|
class TestBroadcast(mlx_tests.MLXTestCase):
|
||
|
|
def test_broadcast_shapes(self):
|
||
|
|
# Basic broadcasting
|
||
|
|
self.assertEqual(mlx.utils.broadcast_shapes((1, 2, 3), (3,)), (1, 2, 3))
|
||
|
|
self.assertEqual(mlx.utils.broadcast_shapes((4, 1, 6), (5, 6)), (4, 5, 6))
|
||
|
|
self.assertEqual(mlx.utils.broadcast_shapes((5, 1, 4), (1, 3, 4)), (5, 3, 4))
|
||
|
|
|
||
|
|
# Multiple arguments
|
||
|
|
self.assertEqual(mlx.utils.broadcast_shapes((1, 1), (1, 8), (7, 1)), (7, 8))
|
||
|
|
self.assertEqual(
|
||
|
|
mlx.utils.broadcast_shapes((6, 1, 5), (1, 7, 1), (6, 7, 5)), (6, 7, 5)
|
||
|
|
)
|
||
|
|
|
||
|
|
# Same shapes
|
||
|
|
self.assertEqual(mlx.utils.broadcast_shapes((3, 4, 5), (3, 4, 5)), (3, 4, 5))
|
||
|
|
|
||
|
|
# Single argument
|
||
|
|
self.assertEqual(mlx.utils.broadcast_shapes((2, 3)), (2, 3))
|
||
|
|
|
||
|
|
# Empty shapes
|
||
|
|
self.assertEqual(mlx.utils.broadcast_shapes((), ()), ())
|
||
|
|
self.assertEqual(mlx.utils.broadcast_shapes((), (1,)), (1,))
|
||
|
|
self.assertEqual(mlx.utils.broadcast_shapes((1,), ()), (1,))
|
||
|
|
|
||
|
|
# Broadcasting with zeroes
|
||
|
|
self.assertEqual(mlx.utils.broadcast_shapes((0,), (0,)), (0,))
|
||
|
|
self.assertEqual(mlx.utils.broadcast_shapes((1, 0, 5), (3, 1, 5)), (3, 0, 5))
|
||
|
|
self.assertEqual(mlx.utils.broadcast_shapes((5, 0), (0, 5, 0)), (0, 5, 0))
|
||
|
|
|
||
|
|
# Error cases
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
mlx.utils.broadcast_shapes((3, 4), (4, 3))
|
||
|
|
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
mlx.utils.broadcast_shapes((2, 3, 4), (2, 5, 4))
|
||
|
|
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
mlx.utils.broadcast_shapes()
|