diff --git a/tensorflow_probability/python/math/hypergeometric.py b/tensorflow_probability/python/math/hypergeometric.py index 55fa7bd7ee..4d9e2be5f4 100644 --- a/tensorflow_probability/python/math/hypergeometric.py +++ b/tensorflow_probability/python/math/hypergeometric.py @@ -431,6 +431,7 @@ def _hyp2f1_z_near_one(a, b, c, z): """"Compute 2F1(a, b, c, z) when z is near 1.""" with tf.name_scope('hyp2f1_z_near_one'): dtype = dtype_util.common_dtype([a, b, c, z], tf.float32) + numpy_dtype = dtype_util.as_numpy_dtype(dtype) a = tf.convert_to_tensor(a, dtype=dtype) b = tf.convert_to_tensor(b, dtype=dtype) c = tf.convert_to_tensor(c, dtype=dtype) @@ -441,6 +442,24 @@ def _hyp2f1_z_near_one(a, b, c, z): d = c - a - b + # The connection formula below has a removable singularity whenever `d` + # is an integer: lgamma(d) has a pole for d <= 0, and lgamma(-d) has a + # pole for d >= 0 (one of the two always fires at an integer d), even + # though 2F1 itself is finite and analytic there. The exact value at + # such a point is given by a distinct closed form with an extra + # log(1 - z) term and digamma corrections (DLMF 15.8.8, 15.8.9), but + # rather than special-case that formula here, we regularize by nudging + # `c` (and thus `d`) off the exact degenerate point. Since 2F1 is + # analytic in `c` away from its own poles, this recovers the true + # limit at a first-order-in-`epsilon` rate; `epsilon = sqrt(machine + # epsilon)` is the standard step size balancing truncation against + # rounding error in a one-sided numerical derivative, and empirically + # matches mpmath to ~1e-8 relative error in float64 (~1e-4 in float32). + epsilon = numpy_dtype(np.sqrt(np.finfo(numpy_dtype).eps)) + is_degenerate = tf.math.equal(d, tf.math.round(d)) + c = tf.where(is_degenerate, c + epsilon, c) + d = tf.where(is_degenerate, d + epsilon, d) + # TODO(b/171982819): When tfp.math.log_gamma_difference and tfp.math.lbeta # support negative parameters, use them here for greater accuracy. log_first_coefficient = (tf.math.lgamma(c) + tf.math.lgamma(d) - diff --git a/tensorflow_probability/python/math/hypergeometric_test.py b/tensorflow_probability/python/math/hypergeometric_test.py index a494f56bff..ecd2eff408 100644 --- a/tensorflow_probability/python/math/hypergeometric_test.py +++ b/tensorflow_probability/python/math/hypergeometric_test.py @@ -197,6 +197,46 @@ def testHyp2F1ParamsLargerCSmaller(self, dtype, rtol): c = self.GenParam(20., 50., dtype, seed_stream()) self.VerifyHyp2F1(dtype, rtol, a, b, c, z_lower=-1., z_upper=1.) + @parameterized.named_parameters( + ("float32", np.float32, 2e-3), + ("float64", np.float64, 1e-7)) + def testHyp2F1DegenerateCMinusAMinusB(self, dtype, rtol): + # Regression test for https://github.com/tensorflow/probability/issues/2001 + # When c - a - b is exactly an integer, the z-near-1 connection formula + # has a removable singularity: lgamma(d) has a pole for d <= 0, and + # lgamma(-d) has a pole for d >= 0, so one of the two always fires, + # producing nan even though 2F1 itself is finite there. Cover a spread + # of m = c - a - b (positive, negative, and the m = 0 boundary). + H = dtype(1.) + a = dtype(1.) + b = dtype(0.5 - H) + c = dtype(H + 1.5) # c - a - b == 2. exactly + z = dtype(0.9901961) + hyp2f1 = self.evaluate(hypergeometric.hyp2f1_small_argument(a, b, c, z)) + self.assertAllClose(hyp2f1, 0.753603006025111, rtol=rtol) + + # Same degenerate (a, b, c) with a different z: the exact repro from + # https://github.com/tensorflow/probability/issues/1861, which is the + # same underlying bug as #2001 above. + z_1861 = dtype(0.9) + hyp2f1_1861 = self.evaluate( + hypergeometric.hyp2f1_small_argument(a, b, c, z_1861)) + self.assertAllClose(hyp2f1_1861, 0.7836799547024426, rtol=rtol) + + for a_val, b_val, c_val, z_val in [ + (0.5, 0.5, 1.0, 0.99), # m == 0 (boundary between the two poles) + (3.0, 2.0, 6.0, 0.999), # m == 1 + (2.3, 1.5, 6.8, 0.95), # m == 3 + (2.0, 1.0, 2.0, 0.95), # m == -1 (c - a - b < 0) + (4.5, 1.7, 3.2, 0.92), # m == -3 + ]: + a_t, b_t, c_t, z_t = (dtype(a_val), dtype(b_val), dtype(c_val), + dtype(z_val)) + hyp2f1 = self.evaluate( + hypergeometric.hyp2f1_small_argument(a_t, b_t, c_t, z_t)) + expected = scipy_special.hyp2f1(a_val, b_val, c_val, z_val) + self.assertAllClose(hyp2f1, expected, rtol=rtol) + @test_util.numpy_disable_gradient_test @test_util.jax_disable_test_missing_functionality( "Gradients not supported in JAX.")