Bug fix in metal binary kernel dispatch for large arrays (#125)

* bug fix * format
2025-10-18 15:28:16 +08:00 · 2023-12-10 16:12:31 -08:00
parent 0cfbfc9904
commit 71d1fff90a
3 changed files with 10 additions and 4 deletions
--- a/python/tests/test_ops.py
+++ b/python/tests/test_ops.py
@@ -1305,6 +1305,11 @@ class TestOps(mlx_tests.MLXTestCase):
                            d_np = np.take(b_mx, np.arange(kth), axis=axis)
                            self.assertTrue(np.all(d_np <= c_mx))

+    def test_large_binary(self):
+        a = mx.ones([1000, 2147484], mx.int8)
+        b = mx.ones([2147484], mx.int8)
+        self.assertEqual((a + b)[0, 0].item(), 2)
+

 if __name__ == "__main__":
    unittest.main()