Revert pull request #329 to fix numerical instabilities.

PiperOrigin-RevId: 486636531
google-deepmind · Nov 7, 2022 · 9dcfa4a · 9dcfa4a
1 parent 509c706
commit 9dcfa4a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/optax/_src/transform.py b/optax/_src/transform.py
@@ -109,8 +109,9 @@ def orderth_norm(g):
 
 def bias_correction(moment, decay, count):
   """Perform bias correction. This becomes a no-op as count goes to infinity."""
+  bias_correction_ = 1 - decay**count
   return jax.tree_util.tree_map(
-      lambda t: t / (1 - decay ** count.astype(t.dtype)), moment)
+      lambda t: t / bias_correction_.astype(t.dtype), moment)
 
 
 def _reject_complex(params):