Fix #23133: erfinv returns Inf for largest float32 value below 1.0#23207
Fix #23133: erfinv returns Inf for largest float32 value below 1.0#23207utsab345 wants to merge 2 commits into
Conversation
…ers in closures
Replace with
in torchtree_impl.py and optree_impl.py. The original lambda captured
(which holds GPU tensors) by reference. When PyTorch's
tree_flatten internally creates a self-referential helper closure that
captures this lambda, it forms a reference cycle:
tree_flatten.helper -> is_leaf lambda -> structure (GPU tensors)
^ /
+---------- captures self --------------+
Python's reference counting cannot break cycles, so GPU tensors remain
allocated until the cyclic GC runs, causing linear OOM in loops.
By capturing id(structure) (an integer) instead, GPU tensors are
excluded from the cycle and can be freed promptly.
Fixes: keras-team#23110
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces clamping in the TensorFlow backend's erfinv implementation to prevent numerical overflow for inputs close to ±1.0, and updates tree traversal implementations to use object IDs (id(x)) instead of identity comparisons (is not) to identify the root structure. The review feedback correctly identifies that unconditional clamping in erfinv breaks mathematical correctness for inputs equal to ±1.0 (which should return ±inf) and inputs outside [-1.0, 1.0] (which should return NaN). The reviewer suggested a robust fix using tf.where to only clamp values strictly within the open interval (-1.0, 1.0).
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if dtype in ["bfloat16", "float16"]: | ||
| return tf.cast( | ||
| tf.math.erfinv(tf.cast(x, tf.float32)), | ||
| dtype, | ||
| ) | ||
| x = tf.cast(x, tf.float32) | ||
| epsilon = tf.constant(1e-7, dtype=tf.float32) | ||
| x = tf.clip_by_value(x, -1.0 + epsilon, 1.0 - epsilon) | ||
| return tf.cast(tf.math.erfinv(x), dtype) | ||
| if dtype == "float32": | ||
| epsilon = tf.constant(1e-7, dtype=tf.float32) | ||
| x = tf.clip_by_value(x, -1.0 + epsilon, 1.0 - epsilon) | ||
| elif dtype == "float64": | ||
| epsilon = tf.constant(1e-15, dtype=tf.float64) | ||
| x = tf.clip_by_value(x, -1.0 + epsilon, 1.0 - epsilon) |
There was a problem hiding this comment.
Unconditionally clamping x to [-1.0 + epsilon, 1.0 - epsilon] introduces correctness bugs for edge cases:
- Inputs equal to
1.0or-1.0: Mathematically,erfinv(1.0) = infanderfinv(-1.0) = -inf. With unconditional clamping, they will return a finite value (approx3.83or5.6) instead ofinf/-inf. - Inputs outside
[-1.0, 1.0]: Mathematically,erfinv(x)is undefined forabs(x) > 1.0and should returnNaN. With unconditional clamping,erfinv(1.5)will return a finite value instead ofNaN.
To fix this, we should only apply the clamping to values strictly within (-1.0, 1.0) (i.e., where abs(x) < 1.0), using tf.where to preserve the original values for abs(x) >= 1.0.
| if dtype in ["bfloat16", "float16"]: | |
| return tf.cast( | |
| tf.math.erfinv(tf.cast(x, tf.float32)), | |
| dtype, | |
| ) | |
| x = tf.cast(x, tf.float32) | |
| epsilon = tf.constant(1e-7, dtype=tf.float32) | |
| x = tf.clip_by_value(x, -1.0 + epsilon, 1.0 - epsilon) | |
| return tf.cast(tf.math.erfinv(x), dtype) | |
| if dtype == "float32": | |
| epsilon = tf.constant(1e-7, dtype=tf.float32) | |
| x = tf.clip_by_value(x, -1.0 + epsilon, 1.0 - epsilon) | |
| elif dtype == "float64": | |
| epsilon = tf.constant(1e-15, dtype=tf.float64) | |
| x = tf.clip_by_value(x, -1.0 + epsilon, 1.0 - epsilon) | |
| if dtype in ["bfloat16", "float16"]: | |
| x_f32 = tf.cast(x, tf.float32) | |
| epsilon = tf.constant(1e-7, dtype=tf.float32) | |
| clamped_x = tf.clip_by_value(x_f32, -1.0 + epsilon, 1.0 - epsilon) | |
| x_f32 = tf.where(tf.abs(x_f32) < 1.0, clamped_x, x_f32) | |
| return tf.cast(tf.math.erfinv(x_f32), dtype) | |
| if dtype == "float32": | |
| epsilon = tf.constant(1e-7, dtype=tf.float32) | |
| clamped_x = tf.clip_by_value(x, -1.0 + epsilon, 1.0 - epsilon) | |
| x = tf.where(tf.abs(x) < 1.0, clamped_x, x) | |
| elif dtype == "float64": | |
| epsilon = tf.constant(1e-15, dtype=tf.float64) | |
| clamped_x = tf.clip_by_value(x, -1.0 + epsilon, 1.0 - epsilon) | |
| x = tf.where(tf.abs(x) < 1.0, clamped_x, x) |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #23207 +/- ##
==========================================
+ Coverage 80.39% 84.04% +3.65%
==========================================
Files 465 465
Lines 69188 69216 +28
Branches 11379 11381 +2
==========================================
+ Hits 55622 58175 +2553
+ Misses 10742 8115 -2627
- Partials 2824 2926 +102
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
utsab345
force-pushed
the
fix/erfinv-boundary
branch
4 times, most recently
from
13020a9 to
f4db822
Compare
utsab345
force-pushed
the
fix/erfinv-boundary
branch
2 times, most recently
from
86acd1c to
bea7c7d
Compare
tf.math.erfinv returns inf for float32 inputs extremely close to 1.0 (e.g. np.nextafter(float32(1.0), 0.0)) despite the input being strictly less than 1.0 where the true erfinv is finite (~3.8325). The fix clamps inputs to [-1.0 + epsilon, 1.0 - epsilon] before passing to tf.math.erfinv: - float32/bfloat16/float16: epsilon = 1e-7 - float64: epsilon = 1e-15 Fixes: keras-team#23133
utsab345
force-pushed
the
fix/erfinv-boundary
branch
from
bea7c7d to
6ced9de
Compare
4 participants
Description
Fixes #23133 (and supersedes #23139).
Fixes
keras.ops.erfinvreturninginffor float32 inputs extremely close to 1.0 where the true inverse error function is finite.Root cause
tf.math.erfinvreturnsinffor inputs likenp.nextafter(np.float32(1.0), np.float32(0.0))(=0.99999994) despite the true mathematical value being finite (~3.8325). This is due to insufficient numerical precision in TF's implementation near the boundary.Fix
Clamp inputs to
[-1.0 + epsilon, 1.0 - epsilon]before passing totf.math.erfinv:Other backends (JAX, NumPy/SciPy, PyTorch) are unaffected.
Contributor Agreement