From 078cf07cd0209aa88205569c3d476eaf0d029e3b Mon Sep 17 00:00:00 2001 From: DiadkoShmek <207617272+DiadkoShmek@users.noreply.github.com> Date: Thu, 5 Nov 2026 02:26:48 +0200 Subject: [PATCH] Restore atan_difference gradients at negative-one products --- .../distributions/truncated_cauchy_test.py | 15 +++++++++ tensorflow_probability/python/math/special.py | 33 ++++++++++++------- .../python/math/special_test.py | 28 ++++++++++++++++ 3 files changed, 64 insertions(+), 12 deletions(-) diff --git a/tensorflow_probability/python/distributions/truncated_cauchy_test.py b/tensorflow_probability/python/distributions/truncated_cauchy_test.py index 572417ec40..e00eff720c 100644 --- a/tensorflow_probability/python/distributions/truncated_cauchy_test.py +++ b/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. """ diff --git a/tensorflow_probability/python/math/special.py b/tensorflow_probability/python/math/special.py index 9d0c472d6a..a6bece1577 100644 --- a/tensorflow_probability/python/math/special.py +++ b/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]] diff --git a/tensorflow_probability/python/math/special_test.py b/tensorflow_probability/python/math/special_test.py index 675bcf2a44..b847809d15 100644 --- a/tensorflow_probability/python/math/special_test.py +++ b/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(