Improve numerical stability of tfb.Softplus forward and its gradient.

PiperOrigin-RevId: 311035111

Author: SiegeLordEx

tensorflow_probability/python/bijectors/BUILD

@@ -1521,12 +1521,16 @@ multi_substrate_py_test(
     name = "softplus_test",
     size = "small",
     srcs = ["softplus_test.py"],
+    jax_size = "medium",
     deps = [
         ":bijector_test_util",
         ":bijectors",
+        # absl/testing:parameterized dep,
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/internal:test_util",
+        "//tensorflow_probability/python/math",
+        # tensorflow/compiler/jit dep,
     ],
 )
 

tensorflow_probability/python/bijectors/softplus.py

@@ -33,6 +33,31 @@
 ]
 
 
+JAX_MODE = False  # Overwritten by rewrite script.
+
+
+# TODO(b/155501444): Remove this when tf.nn.softplus is fixed.
+if JAX_MODE:
+  _stable_grad_softplus = tf.nn.softplus
+else:
+
+  @tf.custom_gradient
+  def _stable_grad_softplus(x):
+    """A (more) numerically stable softplus than `tf.nn.softplus`."""
+    x = tf.convert_to_tensor(x)
+    if x.dtype == tf.float64:
+      cutoff = -20
+    else:
+      cutoff = -9
+
+    y = tf.where(x < cutoff, tf.math.log1p(tf.exp(x)), tf.nn.softplus(x))
+
+    def grad_fn(dy):
+      return dy * tf.where(x < cutoff, tf.exp(x), tf.nn.sigmoid(x))
+
+    return y, grad_fn
+
+
 class Softplus(bijector.Bijector):
   """Bijector which computes `Y = g(X) = Log[1 + exp(X)]`.
 
@@ -101,9 +126,9 @@ def _is_increasing(cls):
 
   def _forward(self, x):
     if self.hinge_softness is None:
-      return tf.math.softplus(x)
+      return _stable_grad_softplus(x)
     hinge_softness = tf.cast(self.hinge_softness, x.dtype)
-    return hinge_softness * tf.math.softplus(x / hinge_softness)
+    return hinge_softness * _stable_grad_softplus(x / hinge_softness)
 
   def _inverse(self, y):
     if self.hinge_softness is None:

tensorflow_probability/python/bijectors/softplus_test.py

@@ -20,9 +20,11 @@
 
 # Dependency imports
 
+from absl.testing import parameterized
 import numpy as np
 import tensorflow.compat.v2 as tf
 from tensorflow_probability.python import bijectors as tfb
+from tensorflow_probability.python import math as tfp_math
 from tensorflow_probability.python.bijectors import bijector_test_util
 from tensorflow_probability.python.internal import test_util
 
@@ -149,6 +151,25 @@ def testVariableHingeSoftness(self):
       with tf.control_dependencies([hinge_softness.assign(0.)]):
         self.evaluate(b.forward(0.5))
 
+  @parameterized.named_parameters(
+      ('32bitGraph', np.float32, False),
+      ('64bitGraph', np.float64, False),
+      ('32bitXLA', np.float32, True),
+      ('64bitXLA', np.float64, True),
+  )
+  @test_util.numpy_disable_gradient_test
+  def testLeftTailGrad(self, dtype, do_compile):
+    x = np.linspace(-50., -8., 1000).astype(dtype)
+
+    @tf.function(autograph=False, experimental_compile=do_compile)
+    def fn(x):
+      return tf.math.log(tfb.Softplus().forward(x))
+
+    _, grad = tfp_math.value_and_gradient(fn, x)
+
+    true_grad = 1 / (1 + np.exp(-x)) / np.log1p(np.exp(x))
+    self.assertAllClose(true_grad, self.evaluate(grad), atol=1e-3)
+
 
 if __name__ == '__main__':
   tf.test.main()