zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-06 12:09:43 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

split logsumexp

Browse Source
This commit is contained in:
Awni Hannun
2025-05-06 17:10:14 -07:00
parent 5a1a5d5ed1
commit 7c99acb799
3 changed files with 42 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/tests/test_ops.py
View File

@@ -760,6 +760,10 @@ class TestOps(mlx_tests.MLXTestCase):
x = mx.broadcast_to(mx.random.uniform(shape=(2, 1, 8)), (2, 2, 8))
self.assertTrue(mx.allclose(mx.logsumexp(x), logsumexp(x)))
# Even larger
x = mx.random.uniform(shape=(4 * 4096 + 3,))
self.assertTrue(mx.allclose(mx.logsumexp(x), logsumexp(x)))
def test_mean(self):
x = mx.array(
[