Add GPU support for uint64/int64 reductions (#569)

2025-09-05 16:13:52 +08:00 · 2024-01-31 11:18:04 -08:00
parent bad67fec37
commit fcc5ac1c64
5 changed files with 654 additions and 192 deletions
--- a/python/tests/test_reduce.py
+++ b/python/tests/test_reduce.py
@@ -55,6 +55,8 @@ class TestReduce(mlx_tests.MLXTestCase):
            "uint8",
            "uint16",
            "uint32",
+            "int64",
+            "uint64",
        ]
        float_dtypes = ["float32"]