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

fix simd erf_inv (#1896)

Browse Source
This commit is contained in:
Awni Hannun 2025-02-24 13:57:47 -08:00 committed by GitHub
parent 8ff84b5c43
commit 2d0f384b6f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mlx/backend/cpu/simd/math.h
View File

@ -186,7 +186,7 @@ Simd<T, N> erfinv(Simd<T, N> a_) {
return a * rhs(t);
}
} else {
return a * select(t > thresh, lhs(t), rhs(t));
return a * select(abs(t) > thresh, lhs(t), rhs(t));
}
}

1
mlx/backend/metal/kernels/CMakeLists.txt
View File

@ -4,6 +4,7 @@ set(BASE_HEADERS
bf16_math.h
complex.h
defines.h
erf.h
expm1f.h
utils.h)

4
python/tests/test_ops.py
View File

@ -898,6 +898,10 @@ class TestOps(mlx_tests.MLXTestCase):
).astype(np.float32)
self.assertTrue(np.allclose(mx.erfinv(x), expected, equal_nan=True))
result = mx.erfinv(mx.array([0.9999999403953552] * 8))
expected = mx.array([3.8325066566467285] * 8)
self.assertTrue(mx.allclose(result, expected))
def test_sin(self):
a = mx.array(
[0, math.pi / 4, math.pi / 2, math.pi, 3 * math.pi / 4, 2 * math.pi]