fix simd erf_inv (#1896)

2025-09-18 10:26:56 +08:00 · 2025-02-24 13:57:47 -08:00
parent 8ff84b5c43
commit 2d0f384b6f
3 changed files with 6 additions and 1 deletions
--- a/mlx/backend/cpu/simd/math.h
+++ b/mlx/backend/cpu/simd/math.h
@@ -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));
  }
 }

--- a/mlx/backend/metal/kernels/CMakeLists.txt
+++ b/mlx/backend/metal/kernels/CMakeLists.txt
@@ -4,6 +4,7 @@ set(BASE_HEADERS
    bf16_math.h
    complex.h
    defines.h
+    erf.h
    expm1f.h
    utils.h)