Fix gradient for softplus at 0 in the TFP JAX backend.

PiperOrigin-RevId: 286582193
davmre · Dec 20, 2019 · b0d44b1 · b0d44b1
1 parent 9e8e380
commit b0d44b1
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/tensorflow_probability/python/internal/backend/numpy/BUILD b/tensorflow_probability/python/internal/backend/numpy/BUILD
@@ -262,7 +262,9 @@ py_library(
     srcs = ["random_generators.py"],
     deps = [
         ":_utils",
+        ":numpy_array",
         ":numpy_math",
+        ":ops",
         # numpy dep,
         # tensorflow dep,
     ],

diff --git a/tensorflow_probability/python/internal/backend/numpy/numpy_math.py b/tensorflow_probability/python/internal/backend/numpy/numpy_math.py
@@ -27,6 +27,7 @@
 
 from tensorflow_probability.python.internal.backend.numpy import _utils as utils
 from tensorflow_probability.python.internal.backend.numpy.numpy_array import _reverse
+from tensorflow_probability.python.internal.backend.numpy.ops import _custom_gradient
 
 scipy_special = utils.try_import('scipy.special')
 
@@ -730,9 +731,17 @@ def _multiply_no_nan(x, y, name=None):  # pylint: disable=unused-argument
     tf.math.softmax,
     _softmax)
 
+
+@_custom_gradient
+def _softplus(x, name=None):  # pylint: disable=unused-argument
+  def grad(dy):
+    return dy * scipy_special.expit(x)
+  # TODO(b/146563881): Investigate improving numerical accuracy here.
+  return np.log1p(np.exp(-np.abs(x))) + np.maximum(x, 0.), grad
+
 softplus = utils.copy_docstring(
     tf.math.softplus,
-    lambda x, name=None: np.log1p(np.exp(-np.abs(x))) + np.maximum(x, 0.))
+    _softplus)
 
 softsign = utils.copy_docstring(
     tf.math.softsign,