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

Precise sigmoid (#2659)

Browse Source 
* bump patch

* Sigmoid matches PyTorch and is more precise on tails
This commit is contained in:
Awni Hannun
2025-10-10 10:05:23 -07:00
committed by GitHub
parent 380aeb58ae
commit 630350ad3e
4 changed files with 12 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
mlx/backend/cpu/unary_ops.h
View File

@@ -77,7 +77,8 @@ struct Real {
struct Sigmoid { struct Sigmoid {
template <int N, typename T> template <int N, typename T>
Simd<T, N> operator()(Simd<T, N> x) { Simd<T, N> operator()(Simd<T, N> x) {
return 1.0f / (1.0f + simd::exp(-x)); auto y = 1.0f / (1.0f + simd::exp(simd::abs(x)));
return simd::select(x < Simd<T, N>{0}, y, Simd<T, N>{1} - y);
} }
SINGLE() SINGLE()
}; };

4
mlx/backend/cuda/device/unary_ops.cuh
View File

@@ -257,8 +257,8 @@ struct Round {
struct Sigmoid { struct Sigmoid {
template <typename T> template <typename T>
__device__ T operator()(T x) { __device__ T operator()(T x) {
T y = 1 / (1 + exp(-abs(x))); T y = 1 / (1 + exp(abs(x)));
return (x < 0) ? 1 - y : y; return (x < 0) ? y : 1 - y;
} }
}; };

4
mlx/backend/metal/kernels/unary_ops.h
View File

@@ -309,8 +309,8 @@ struct Round {
struct Sigmoid { struct Sigmoid {
template <typename T> template <typename T>
T operator()(T x) { T operator()(T x) {
auto y = 1 / (1 + metal::exp(-metal::abs(x))); auto y = 1 / (1 + metal::exp(metal::abs(x)));
return (x < 0) ? 1 - y : y; return (x < 0) ? y : 1 - y;
} }
}; };

6
python/tests/test_ops.py
View File

@@ -1041,6 +1041,12 @@ class TestOps(mlx_tests.MLXTestCase):
expected = 1 / (1 + np.exp(-a, dtype=np.float32)) expected = 1 / (1 + np.exp(-a, dtype=np.float32))
self.assertTrue(np.allclose(result, expected)) self.assertTrue(np.allclose(result, expected))
# Low precision
a = mx.array(-8.0).astype(mx.float16)
self.assertNotEqual(mx.sigmoid(a).item(), 0.0)
a = mx.array(8.0).astype(mx.float16)
self.assertNotEqual(mx.sigmoid(a).item(), 1.0)
def test_allclose(self): def test_allclose(self):
a = mx.array(1.0) a = mx.array(1.0)
b = mx.array(1.0) b = mx.array(1.0)