Precise sigmoid (#2659)

* bump patch * Sigmoid matches PyTorch and is more precise on tails
2025-12-16 01:49:05 +08:00 · 2025-10-10 10:05:23 -07:00
parent 380aeb58ae
commit 630350ad3e
4 changed files with 12 additions and 5 deletions
--- a/mlx/backend/cpu/unary_ops.h
+++ b/mlx/backend/cpu/unary_ops.h
@@ -77,7 +77,8 @@ struct Real {
 struct Sigmoid {
  template <int N, typename T>
  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()
 };