zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-31 16:21:27 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

consistently handle all -inf in softmax (#1470)

Browse Source
This commit is contained in:
Awni Hannun
2024-10-08 09:54:02 -07:00
committed by GitHub
parent 3274c6a087
commit 1fa0d20a30
3 changed files with 13 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/tests/test_ops.py
View File

@@ -1567,6 +1567,11 @@ class TestOps(mlx_tests.MLXTestCase):
out = mx.softmax(a, axis=-1, precise=True)
self.assertTrue(mx.allclose(out_expect, out))
# All Infs give NaNs
for n in [127, 128, 129]:
x = mx.full((n,), vals=-float("inf"))
self.assertTrue(mx.all(mx.isnan(mx.softmax(x))))
def test_concatenate(self):
a_npy = np.random.randn(32, 32, 32)
b_npy = np.random.randn(32, 32, 32)