Refactor reductions and fix scatter atomics for large sizes (#1300)

Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>
This commit is contained in:
Awni Hannun
2024-08-22 16:03:31 -07:00
committed by GitHub
parent f9e00efe31
commit 98b6ce3460
18 changed files with 1584 additions and 1235 deletions

View File

@@ -1763,7 +1763,7 @@ class TestOps(mlx_tests.MLXTestCase):
mat_t = mat.astype(t)
out = mx.cumsum(a_t, axis=-1)
expected = (mat_t * a_t[:, None, :]).sum(axis=-1)
self.assertTrue(mx.allclose(out, expected, rtol=1e-2, atol=1e-3))
self.assertTrue(mx.allclose(out, expected, rtol=0.02, atol=1e-3))
sizes = [1023, 1024, 1025, 2047, 2048, 2049]
for s in sizes:
a = mx.ones((s,), mx.int32)