Add tfp.math.smootherstep, ie:

```python
x = tf.clip_by_value(x, clip_value_min=0., clip_value_max=1.)
y = x**3. * (6. * x**2. - 15. * x + 10.)
```

PiperOrigin-RevId: 270999562

tensorflow_probability/python/math/__init__.py

@@ -26,6 +26,7 @@
 from tensorflow_probability.python.math.generic import log_combinations
 from tensorflow_probability.python.math.generic import reduce_logmeanexp
 from tensorflow_probability.python.math.generic import reduce_weighted_logsumexp
+from tensorflow_probability.python.math.generic import smootherstep
 from tensorflow_probability.python.math.generic import soft_threshold
 from tensorflow_probability.python.math.generic import softplus_inverse
 from tensorflow_probability.python.math.gradient import value_and_gradient
@@ -74,13 +75,14 @@
     'minimize',
     'ode',
     'pinv',
-    'psd_kernels',
     'pivoted_cholesky',
+    'psd_kernels',
     'random_rademacher',
     'random_rayleigh',
     'reduce_logmeanexp',
     'reduce_weighted_logsumexp',
     'secant_root',
+    'smootherstep',
     'soft_threshold',
     'softplus_inverse',
     'sparse_or_dense_matmul',

tensorflow_probability/python/math/generic.py

@@ -33,6 +33,7 @@
     'log_combinations',
     'reduce_logmeanexp',
     'reduce_weighted_logsumexp',
+    'smootherstep',
     'soft_threshold',
     'softplus_inverse',
 ]
@@ -368,3 +369,42 @@ def log_add_exp(x, y, name=None):
     x = tf.convert_to_tensor(x, dtype=dtype, name='x')
     y = tf.convert_to_tensor(y, dtype=dtype, name='y')
     return tf.maximum(x, y) + tf.math.softplus(-abs(x - y))
+
+
+def smootherstep(x, name=None):
+  """Computes a sigmoid-like interpolation function on the unit-interval.
+
+  Equivalent to:
+
+  ```python
+  x = tf.clip_by_value(x, clip_value_min=0., clip_value_max=1.)
+  y = x**3. * (6. * x**2. - 15. * x + 10.)
+  ```
+
+  For more details see [Wikipedia][1].
+
+  Args:
+    x: `float` `Tensor`.
+    name: Python `str` name prefixed to Ops created by this function.
+      Default value: `None` (i.e., `'smootherstep'`).
+
+  Returns:
+    smootherstep: `float` `Tensor` with the same shape and dtype as `x`,
+      representing the value of the smootherstep function.
+
+  #### References
+
+  [1]: "Smoothstep." Wikipedia.
+       https://en.wikipedia.org/wiki/Smoothstep#Variations
+  """
+  with tf.name_scope(name or 'smootherstep'):
+    x = tf.clip_by_value(x, clip_value_min=0., clip_value_max=1.)
+    # Note: Grappler will rewrite:
+    #   x**2, x**3
+    # as:
+    #   x2 = tf.square(x)
+    #   x3 = tf.square(x) * x
+    # and common subexpression elimination (CSE) will produce:
+    #   x2 = tf.square(x)
+    #   x3 = x2 * x
+    return x**3. * (6. * x**2. - 15. * x + 10.)

tensorflow_probability/python/math/generic_test.py

@@ -340,17 +340,22 @@ def test_big(self):
 
 
 @test_util.run_all_in_graph_and_eager_modes
-class ReduceLogMeanExp(test_case.TestCase):
-
-  def test_vector_axis_and_keepdims(self):
-    log_probs = tf.math.log(tf.random.uniform([10, 3, 4], seed=41))
-    expected = tf.math.log(
-        tf.reduce_mean(
-            tf.math.exp(log_probs),
-            axis=[1, 2],
-            keepdims=True))
-    actual = tfp.math.reduce_logmeanexp(log_probs, axis=[1, 2], keepdims=True)
-    self.assertAllClose(*self.evaluate([expected, actual]), rtol=1e-5, atol=0.)
+class Smootherstep(test_case.TestCase):
+
+  def test_value_vector(self):
+    x = tf.constant([-np.inf, -20., 0., 0.5, 1., 20., np.inf])
+    y, _ = tfp.math.value_and_gradient(tfp.math.smootherstep, x)
+    self.assertAllEqual([7], y.shape)
+    y_ = self.evaluate(y)
+    self.assertAllClose([0., 0., 0., 0.5, 1., 1., 1.], y_, atol=1e-5, rtol=1e-5)
+
+  def test_gradient_matrix(self):
+    x = tf.constant([[-np.inf, -20., 0., 0.5],
+                     [np.inf, 20., 1., 0.5]])
+    _, g = tfp.math.value_and_gradient(tfp.math.smootherstep, x)
+    self.assertAllEqual([2, 4], g.shape)
+    g_ = self.evaluate(g)
+    self.assertAllClose([[0., 0., 0., 1.875]] * 2, g_, atol=1e-5, rtol=1e-5)
 
 
 if __name__ == '__main__':