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

Fix overflow / underflow handling for expm1f (#1278)

Browse Source 
* Fix overflow / underflow handling for expm1f

* update tests
This commit is contained in:
Tim Gymnich 2024-07-23 16:29:06 +02:00 committed by GitHub
parent 1fba87b0df
commit 6307d166eb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
mlx/backend/metal/kernels/expm1f.h
View File

@ -83,6 +83,7 @@ float expm1f(float a) {
r = expm1f_scaled_unchecked(a, 1.0f); r = expm1f_scaled_unchecked(a, 1.0f);
/* handle severe overflow and underflow */ /* handle severe overflow and underflow */
if (abs(a - 1.0f) > 88.0f) { if (abs(a - 1.0f) > 88.0f) {
r = pow(2, a);
r = fma(r, r, -1.0f); r = fma(r, r, -1.0f);
} }
return r; return r;

2
python/tests/test_ops.py
View File

@ -845,7 +845,7 @@ class TestOps(mlx_tests.MLXTestCase):
self.assertTrue(np.allclose(result, expected)) self.assertTrue(np.allclose(result, expected))
def test_expm1(self): def test_expm1(self):
a = mx.array([0, 0.5, -0.5, 5]) a = mx.array([-88, -87, 0, 0.5, -0.5, 5, 87, 88, 89, 90])
result = mx.expm1(a) result = mx.expm1(a)
expected = np.expm1(a, dtype=np.float32) expected = np.expm1(a, dtype=np.float32)