zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-05 07:34:42 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Reduce specializations (#1607)

Browse Source 
* start of reduce specializations

* fix all reduce

* fix many dims

* fix

* non-jit tests clear

* cleanup instantiations

* cpu merges

* change dim specializations

* optimize

* fix jit

* fix jit

* use higher precision for integer sum+prod

* fixes
This commit is contained in:
Awni Hannun
2024-11-21 19:53:00 -08:00
committed by GitHub
parent dcca0d7477
commit 0c5eea226b
14 changed files with 733 additions and 406 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
python/tests/test_reduce.py
View File

@@ -131,6 +131,28 @@ class TestReduce(mlx_tests.MLXTestCase):
mxsum = y.sum().item()
self.assertEqual(npsum, mxsum)
def test_many_reduction_axes(self):
def check(x, axes):
expected = x
for ax in axes:
expected = mx.sum(expected, axis=ax, keepdims=True)
out = mx.sum(x, axis=axes, keepdims=True)
self.assertTrue(mx.array_equal(out, expected))
x = mx.random.randint(0, 10, shape=(4, 4, 4, 4, 4))
check(x, (0, 2, 4))
x = mx.random.randint(0, 10, shape=(4, 4, 4, 4, 4, 4, 4))
check(x, (0, 2, 4, 6))
x = mx.random.randint(0, 10, shape=(4, 4, 4, 4, 4, 4, 4, 4, 4))
check(x, (0, 2, 4, 6, 8))
x = mx.random.randint(0, 10, shape=(4, 4, 4, 4, 4, 4, 4, 4, 4, 128))
x = x.transpose(1, 0, 2, 3, 4, 5, 6, 7, 8, 9)
check(x, (1, 3, 5, 7, 9))
if __name__ == "__main__":
unittest.main(failfast=True)