Fix #23133: erfinv returns Inf for largest float32 value below 1.0

keras/src/backend/tensorflow/math.py

@@ -305,11 +305,23 @@ def erfinv(x):
     x = convert_to_tensor(x)
     dtype = standardize_dtype(x.dtype)
 
+    # Clamp to avoid numerical overflow (inf) for inputs extremely close
+    # to ±1.0 where tf.math.erfinv returns inf despite the input being
+    # strictly less than 1.0 (e.g. np.nextafter(float32(1.0), 0.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)
+        clamped_x = tf.clip_by_value(x, -1.0 + epsilon, 1.0 - epsilon)
+        x = tf.where(tf.abs(x) < 1.0, clamped_x, x)
+        return tf.cast(tf.math.erfinv(x), 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)
 
     return tf.math.erfinv(x)
 

keras/src/tree/optree_impl.py

@@ -46,10 +46,12 @@ def is_nested(structure):
 
 def traverse(func, structure, top_down=True):
     # From https://github.com/google/jax/pull/19695
+    structure_id = id(structure)
+
     def traverse_children():
         children, treedef = optree.tree_flatten(
             structure,
-            is_leaf=lambda x: x is not structure,
+            is_leaf=lambda x: id(x) != structure_id,
             none_is_leaf=True,
             namespace="keras",
         )

keras/src/tree/torchtree_impl.py

@@ -33,10 +33,12 @@ def func(x):
             )
         return None
 
+    structure_id = id(structure)
+
     def traverse_children():
         children, treedef = torch_tree.tree_flatten(
             structure,
-            is_leaf=lambda x: x is not structure,
+            is_leaf=lambda x: id(x) != structure_id,
         )
         if treedef.num_nodes == 1 and treedef.num_leaves == 1:
             return structure
@@ -59,10 +61,13 @@ def is_nested(structure):
 
 
 def traverse(func, structure, top_down=True):
+    structure = _dict_to_ordered_dict(structure)
+    structure_id = id(structure)
+
     def traverse_children():
         children, treedef = torch_tree.tree_flatten(
             structure,
-            is_leaf=lambda x: x is not structure,
+            is_leaf=lambda x: id(x) != structure_id,
         )
         if treedef.num_nodes == 1 and treedef.num_leaves == 1:
             return structure
@@ -72,7 +77,6 @@ def traverse_children():
                 treedef,
             )
 
-    structure = _dict_to_ordered_dict(structure)
     if top_down:
         ret = func(structure)
         if ret is None: