Merge pull request jax-ml#27107 from pearu:pearu/betainc

PiperOrigin-RevId: 738855631

Author: Google-ML-Automation

jax/_src/lax/special.py

@@ -194,12 +194,18 @@ def nth_partial_betainc_numerator(iteration, a, b, x):
     iteration_is_one = eq(iteration_bcast, full_like(iteration_bcast, 1))
     iteration_minus_one = iteration_bcast - full_like(iteration_bcast, 1)
     m = iteration_minus_one // full_like(iteration_minus_one, 2)
+    m_is_zero = eq(m, full_like(m, 0))
     m = convert_element_type(m, dtype)
     one = full_like(a, 1)
     two = full_like(a, 2.0)
     # Partial numerator terms
-    even_numerator = -(a + m) * (a + b + m) * x / (
-        (a + two * m) * (a + two * m + one))
+
+    # When a is close to zero and m == 0, using zero_numerator avoids
+    # inaccuracies when FTZ or DAZ is enabled:
+    zero_numerator = -(a + b) * x / (a + one)
+    even_numerator = select(m_is_zero, zero_numerator,
+                            -(a + m) * (a + b + m) * x / (
+                              (a + two * m) * (a + two * m + one)))
     odd_numerator = m * (b - m) * x / ((a + two * m - one) * (a + two * m))
     one_numerator = full_like(x, 1.0)
     numerator = select(iteration_is_even, even_numerator, odd_numerator)
@@ -210,12 +216,24 @@ def nth_partial_betainc_denominator(iteration, a, b, x):
     return select(eq(iteration_bcast, full_like(iteration_bcast, 0)),
                   full_like(x, 0), full_like(x, 1))
 
+  a_is_zero = bitwise_or(eq(a, full_like(a, 0)), eq(b, full_like(b, float('inf'))))
+  b_is_zero = bitwise_or(eq(b, full_like(b, 0)), eq(a, full_like(a, float('inf'))))
+  x_is_zero = eq(x, full_like(x, 0))
+  x_is_one = eq(x, full_like(x, 1))
+  x_is_not_zero = bitwise_not(x_is_zero)
+  x_is_not_one = bitwise_not(x_is_one)
+  is_nan = bitwise_or(bitwise_or(_isnan(a), _isnan(b)), _isnan(x))
+
+  result_is_zero = bitwise_or(bitwise_and(b_is_zero, x_is_not_one), bitwise_and(a_is_zero, x_is_zero))
+  result_is_one = bitwise_or(bitwise_and(a_is_zero, x_is_not_zero), bitwise_and(b_is_zero, x_is_one))
+
   result_is_nan = bitwise_or(bitwise_or(bitwise_or(
-    le(a, full_like(a, 0)), le(b, full_like(b, 0))),
+    lt(a, full_like(a, 0)), lt(b, full_like(b, 0))),
     lt(x, full_like(x, 0))), gt(x, full_like(x, 1)))
+  result_is_nan = bitwise_or(result_is_nan, bitwise_or(bitwise_and(a_is_zero, b_is_zero), is_nan))
 
-  # The continued fraction will converge rapidly when x < (a+1)/(a+b+2)
-  # as per: http://dlmf.nist.gov/8.17.E23
+  # The continued fraction will converge rapidly when x <
+  # (a+1)/(a+b+2) as per: http://dlmf.nist.gov/8.17.E23.
   #
   # Otherwise, we can rewrite using the symmetry relation as per:
   # http://dlmf.nist.gov/8.17.E4
@@ -234,10 +252,21 @@ def nth_partial_betainc_denominator(iteration, a, b, x):
     inputs=[a, b, x]
   )
 
-  lbeta_ab = lgamma(a) + lgamma(b) - lgamma(a + b)
-  result = continued_fraction * exp(log(x) * a + log1p(-x) * b - lbeta_ab) / a
+  # For very small a and to avoid division by zero, we'll use
+  # a * gamma(a) = gamma(a + 1) -> 1 as a -> 0+.
+  very_small = (dtypes.finfo(dtype).tiny * 2).astype(dtype)
+  lbeta_ab_small_a = lgamma(b) - lgamma(a + b)
+  lbeta_ab = lgamma(a) + lbeta_ab_small_a
+  factor = select(lt(a, full_like(a, very_small)),
+                  exp(log1p(-x) * b - lbeta_ab_small_a),
+                  exp(log(x) * a + log1p(-x) * b - lbeta_ab) / a)
+  result = continued_fraction * factor
+  result = select(converges_rapidly, result, sub(full_like(result, 1), result))
+
+  result = select(result_is_zero, full_like(a, 0), result)
+  result = select(result_is_one, full_like(a, 1), result)
   result = select(result_is_nan, full_like(a, float('nan')), result)
-  return select(converges_rapidly, result, sub(full_like(result, 1), result))
+  return result
 
 class IgammaMode(Enum):
   VALUE = 1

jax/_src/test_util.py

@@ -1522,7 +1522,7 @@ def _CheckAgainstNumpy(self, numpy_reference_op, lax_op, args_maker,
     args = args_maker()
     lax_ans = lax_op(*args)
     numpy_ans = numpy_reference_op(*args)
-    self.assertAllClose(numpy_ans, lax_ans, check_dtypes=check_dtypes,
+    self.assertAllClose(lax_ans, numpy_ans, check_dtypes=check_dtypes,
                         atol=atol or tol, rtol=rtol or tol,
                         canonicalize_dtypes=canonicalize_dtypes)
 

tests/lax_scipy_special_functions_test.py

@@ -288,35 +288,71 @@ def testExpiDisableJit(self):
     self.assertAllClose(result_jit, result_nojit)
 
   def testGammaIncBoundaryValues(self):
-    dtype = jax.numpy.zeros(0).dtype  # default float dtype.
+    dtype = jax.dtypes.canonicalize_dtype(float)
     nan = float('nan')
     inf = float('inf')
     if jtu.parse_version(scipy.__version__) >= (1, 16):
-      samples_slice = slice(None)
+      a_samples = [0, 0, 0, 1, nan,   1, nan,   0,   1, 1, nan]
+      x_samples = [0, 1, 2, 0,   1, nan, nan, inf, inf, -1, inf]
     else:
       # disable samples that contradict with scipy/scipy#22441
-      samples_slice = slice(None, -1)
-    args_maker = lambda: [np.array([0, 0, 0, 1, nan,   1, nan,   0,   1, nan][samples_slice]).astype(dtype),
-                          np.array([0, 1, 2, 0,   1, nan, nan, inf, inf, inf][samples_slice]).astype(dtype)]
+      a_samples = [0, 0, 0, 1, nan,   1, nan,   0,   1, 1]
+      x_samples = [0, 1, 2, 0,   1, nan, nan, inf, inf, -1]
+
+    args_maker = lambda: (np.array(a_samples, dtype=dtype), np.array(x_samples, dtype=dtype))
+
     rtol = 1E-3 if jtu.test_device_matches(["tpu"]) else 1e-5
     self._CheckAgainstNumpy(lsp_special.gammainc, osp_special.gammainc, args_maker, rtol=rtol)
     self._CompileAndCheck(lsp_special.gammainc, args_maker, rtol=rtol)
 
   def testGammaIncCBoundaryValues(self):
-    dtype = jax.numpy.zeros(0).dtype  # default float dtype.
+    dtype = jax.dtypes.canonicalize_dtype(float)
     nan = float('nan')
     inf = float('inf')
     if jtu.parse_version(scipy.__version__) >= (1, 16):
-      samples_slice = slice(None)
+      a_samples = [0, 0, 0, 1, nan,   1, nan,   0,   1, 1, nan]
+      x_samples = [0, 1, 2, 0,   1, nan, nan, inf, inf, -1, inf]
     else:
       # disable samples that contradict with scipy/scipy#22441
-      samples_slice = slice(None, -1)
-    args_maker = lambda: [np.array([0, 0, 0, 1, nan,   1, nan,   0,   1, 1, nan][samples_slice]).astype(dtype),
-                          np.array([0, 1, 2, 0,   1, nan, nan, inf, inf, -1, inf][samples_slice]).astype(dtype)]
+      a_samples = [0, 0, 0, 1, nan,   1, nan,   0,   1, 1]
+      x_samples = [0, 1, 2, 0,   1, nan, nan, inf, inf, -1]
+
+    args_maker = lambda: (np.array(a_samples, dtype=dtype), np.array(x_samples, dtype=dtype))
+
     rtol = 1E-3 if jtu.test_device_matches(["tpu"]) else 1e-5
     self._CheckAgainstNumpy(lsp_special.gammaincc, osp_special.gammaincc, args_maker, rtol=rtol)
     self._CompileAndCheck(lsp_special.gammaincc, args_maker, rtol=rtol)
 
+  def testBetaIncBoundaryValues(self):
+    dtype = jax.dtypes.canonicalize_dtype(float)
+    fi = jax.numpy.finfo(dtype)
+    nan = float('nan')
+    inf = float('inf')
+    tiny = fi.tiny
+    eps = fi.eps
+    if jtu.parse_version(scipy.__version__) >= (1, 16):
+      # TODO(pearu): enable tiny samples when a fix to scipy/scipy#22682
+      # will be available
+      a_samples = [nan, -0.5, inf, 0, eps, 1, tiny][:-1]
+      b_samples = [nan, -0.5, inf, 0, eps, 1, tiny][:-1]
+    elif jtu.parse_version(scipy.__version__) >= (1, 12):
+      # disabled samples that contradict with scipy/scipy#22425
+      a_samples = [nan, -0.5, 0.5]
+      b_samples = [nan, -0.5, 0.5]
+    else:
+      a_samples = [-0.5, 0.5]
+      b_samples = [-0.5, 0.5]
+    x_samples = [nan, -0.5, 0, 0.5, 1, 1.5]
+
+    a_samples = np.array(a_samples, dtype=dtype)
+    b_samples = np.array(b_samples, dtype=dtype)
+    x_samples = np.array(x_samples, dtype=dtype)
+
+    args_maker = lambda: np.meshgrid(a_samples, b_samples, x_samples)
+
+    rtol = 1E-3 if jtu.test_device_matches(["tpu"]) else 5e-5
+    self._CheckAgainstNumpy(osp_special.betainc, lsp_special.betainc, args_maker, rtol=rtol)
+    self._CompileAndCheck(lsp_special.betainc, args_maker, rtol=rtol)
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())