Restore atan_difference gradients at negative-one products

tensorflow_probability/python/distributions/truncated_cauchy_test.py

@@ -21,11 +21,26 @@
 from tensorflow_probability.python.distributions import cauchy
 from tensorflow_probability.python.distributions import truncated_cauchy
 from tensorflow_probability.python.internal import test_util
+from tensorflow_probability.python.math import gradient
 
 
 @test_util.test_all_tf_execution_regimes
 class TruncatedCauchyTest(test_util.TestCase):
 
+  @test_util.numpy_disable_gradient_test
+  @parameterized.parameters(np.float32, np.float64)
+  def testLogProbGradientsAtNegativeOneProduct(self, dtype):
+    def log_prob(loc, scale):
+      return truncated_cauchy.TruncatedCauchy(
+          loc=loc, scale=scale, low=-1., high=1.).log_prob(.25)
+
+    _, (dloc, dscale) = gradient.value_and_gradient(
+        log_prob, [tf.convert_to_tensor(0., dtype),
+                   tf.convert_to_tensor(1., dtype)])
+    dloc, dscale = self.evaluate([dloc, dscale])
+    self.assertAllClose(8. / 17., dloc)
+    self.assertAllClose(-15. / 17. + 2. / np.pi, dscale)
+
   def testBatchSampling(self):
     """Check (empirically) that different parameters in a batch are respected.
     """

tensorflow_probability/python/math/special.py

@@ -81,13 +81,22 @@ def atan_difference(x, y, name=None):
     x = tf.convert_to_tensor(x, dtype=dtype)
     y = tf.convert_to_tensor(y, dtype=dtype)
 
-    difference = tf.math.atan((x - y) / (1 + x * y))
+    product = x * y
+    product_is_negative_one = tf.math.equal(product, -1.)
+    # Keep inactive arithmetic finite when vmap lowers a conditional to selects.
+    denominator = tf.where(
+        product_is_negative_one, tf.ones_like(product), 1. + product)
+    difference = tf.math.atan((x - y) / denominator)
     difference = difference + tf.where(
-        x * y < - 1., np.pi * tf.math.sign(x), 0.)
-    difference = tf.where(
-        tf.math.equal(x * y, -1.), np.pi * tf.math.sign(x) / 2., difference)
+        product < -1., np.pi * tf.math.sign(x), 0.)
 
-    return difference
+    # The extra arctangents are only needed at opposite-sign poles, where their
+    # difference has no cancellation and retains the smooth partial derivatives.
+    return ps.cond(
+        tf.reduce_any(product_is_negative_one),
+        lambda: tf.where(product_is_negative_one,
+                         tf.math.atan(x) - tf.math.atan(y), difference),
+        lambda: difference)
 
 
 # 16-bit (half precision) floating-point dtypes available on current backend.
@@ -99,7 +108,7 @@ def _betainc_naive(a, b, x):
   dtype_orig = dtype_util.common_dtype([a, b, x], tf.float32)
   # We promote bfloat16 and float16 to float32 to make this function consistent
   # with the XLA implementation of betainc.
-  should_promote_dtype = (dtype_orig in _f16bit_dtypes)
+  should_promote_dtype = dtype_orig in _f16bit_dtypes
   dtype = tf.float32 if should_promote_dtype else dtype_orig
 
   a, b, x = [tf.convert_to_tensor(z, dtype=dtype_orig) for z in [a, b, x]]
@@ -356,7 +365,7 @@ def _betainc_der_power_series(a, b, x, dtype, use_power_series):
   # When x >= a / (a + b), we must apply the symmetry relation given here:
   # https://dlmf.nist.gov/8.17.E4
   #   betainc(a, b, x) = 1 - betainc(b, a, 1 - x)
-  use_symmetry_relation = (safe_x >= safe_a / (safe_a + safe_b))
+  use_symmetry_relation = safe_x >= safe_a / (safe_a + safe_b)
   safe_a_orig = safe_a
   safe_a = tf.where(use_symmetry_relation, safe_b, safe_a)
   safe_b = tf.where(use_symmetry_relation, safe_a_orig, safe_b)
@@ -433,7 +442,7 @@ def _betainc_partials(a, b, x):
   dtype_orig = dtype_util.common_dtype([a, b, x], tf.float32)
   # We promote bfloat16 and float16 to float32 to make this function consistent
   # with betainc.
-  should_promote_dtype = (dtype_orig in _f16bit_dtypes)
+  should_promote_dtype = dtype_orig in _f16bit_dtypes
   dtype = tf.float32 if should_promote_dtype else dtype_orig
   numpy_dtype = dtype_util.as_numpy_dtype(dtype)
   zero = tf.constant(0., dtype=dtype)
@@ -645,7 +654,7 @@ def _betaincinv_computation(a, b, y):
   dtype_orig = dtype_util.common_dtype([a, b, y], tf.float32)
   # We promote bfloat16 and float16 to float32 to make this function consistent
   # with betainc.
-  should_promote_dtype = (dtype_orig in _f16bit_dtypes)
+  should_promote_dtype = dtype_orig in _f16bit_dtypes
   dtype = tf.float32 if should_promote_dtype else dtype_orig
   numpy_dtype = dtype_util.as_numpy_dtype(dtype)
   zero = numpy_dtype(0.)
@@ -677,7 +686,7 @@ def _betaincinv_computation(a, b, y):
     use_symmetry_relation = (error_at_half < zero) & a_and_b_are_small & (
         error_at_mean < zero)
   else:
-    use_symmetry_relation = (error_at_half < zero)
+    use_symmetry_relation = error_at_half < zero
 
   a_orig, y_orig = (a, y)
   a = tf.where(use_symmetry_relation, b, a)
@@ -727,7 +736,7 @@ def root_finding_iteration(should_stop, low, high, candidate):
             halley_candidate))
 
     new_delta = candidate - new_candidate
-    new_delta_is_negative = (new_delta < zero)
+    new_delta_is_negative = new_delta < zero
     new_low = tf.where(new_delta_is_negative, candidate, low)
     new_high = tf.where(new_delta_is_negative, high, candidate)
 
@@ -782,7 +791,7 @@ def _betaincinv_partials(a, b, y, return_value=False):
   dtype_orig = dtype_util.common_dtype([a, b, y], tf.float32)
   # We promote bfloat16 and float16 to float32 to make this function consistent
   # with betaincinv.
-  should_promote_dtype = (dtype_orig in _f16bit_dtypes)
+  should_promote_dtype = dtype_orig in _f16bit_dtypes
   dtype = tf.float32 if should_promote_dtype else dtype_orig
 
   a, b, y = [tf.convert_to_tensor(z, dtype=dtype_orig) for z in [a, b, y]]

tensorflow_probability/python/math/special_test.py

@@ -1134,6 +1134,34 @@ def testAtanDifferenceProductIsNegativeOne(self, dtype):
     x_, y_, atan_diff_ = self.evaluate([x, y, special.atan_difference(x, y)])
     self.assertAllClose(np.arctan(x_) - np.arctan(y_), atan_diff_)
 
+  @test_util.numpy_disable_gradient_test
+  @parameterized.product(
+      dtype=[np.float32, np.float64], broadcast=[False, True])
+  def testAtanDifferenceGradientsAtNegativeOneProduct(self, dtype, broadcast):
+    x = np.array([1., 2., -1., -2.], dtype=dtype)
+    y = -1. / x
+    if broadcast:
+      y = y[:, np.newaxis]
+    x_tensor = tf.convert_to_tensor(x)
+    y_tensor = tf.convert_to_tensor(y)
+    value, (dx, dy) = gradient.value_and_gradient(
+        special.atan_difference, [x_tensor, y_tensor])
+    value, dx, dy = self.evaluate([value, dx, dy])
+    multiplicity = 4. if broadcast else 1.
+    self.assertAllClose(np.arctan(x) - np.arctan(y), value)
+    self.assertAllClose(multiplicity / (1. + x ** 2), dx)
+    self.assertAllClose(-multiplicity / (1. + y ** 2), dy)
+
+    # ForwardAccumulator does not work under tf1.Graph (b/197680518).
+    if test_util.TF_MODE and tf.executing_eagerly():
+      with tf.autodiff.ForwardAccumulator(
+          primals=(x_tensor, y_tensor),
+          tangents=(tf.ones_like(x_tensor), tf.zeros_like(y_tensor))) as acc:
+        traced_value = tf.function(special.atan_difference)(x_tensor, y_tensor)
+      self.assertAllClose(
+          np.broadcast_to(1. / (1. + x ** 2), value.shape),
+          self.evaluate(acc.jvp(traced_value)))
+
   @parameterized.parameters(np.float32, np.float64)
   def testErfcinvPreservesDtype(self, dtype):
     x = self.evaluate(