Repository: google/jax

Issue:
`betaln` is wildly inaccurate for large values
### Description

`jax.scipy.special.betaln` becomes very inaccurate for large inputs.

<img width="608" alt="Screen Shot 2022-11-15 at 4 54 10 PM" src="https://user-images.githubusercontent.com/391217/202057219-814b0ca5-ea8c-41c0-8f5c-748e19fa8fe0.png">

 The image is comparing jax, scipy, and a jax version I made that parallels scipy's version (at the expense of duplicated work, due to the evaluation of both branches).

The more accurate jax version can be found [here](https://gist.github.com/imh/b6e3e7e2ea3224e365360cea85578240). I would be happy to open up a PR for it, if someone could answer the following:

- If this went in `jax._src.scipy.special`, where should I import `jnp.where` from to avoid circular dependencies?
- Since I just ported (two branches of) of scipy's version of `betaln`, I expect there may be be license issues including it in jax. Is that the case? Scipy's license is linked in the ported version.

### What jax/jaxlib version are you using?

0.3.23

### Which accelerator(s) are you using?

CPU

### Additional system info

MacOS 12.6

### NVIDIA GPU info

_No response_


Discussion:
We welcome PRs!

* The circularity with `jnp.where` usually isn't a problem because `jax.scipy` ends up imported after `jax.numpy`. Try it?
* In general if the code is copied or derived from code elsewhere, it must go in the `third_party` subdirectory (https://github.com/google/jax/tree/main/jax/_src/third_party) of `jax._src` since it is under a different license and we need to comply with that license. That's probably fine for SciPy code, and indeed there's already a little scipy-derived code in there you can follow as an example.


Resolving patch:
diff --git a/jax/_src/scipy/special.py b/jax/_src/scipy/special.py
--- a/jax/_src/scipy/special.py
+++ b/jax/_src/scipy/special.py
@@ -29,6 +29,7 @@
 from jax._src.numpy.lax_numpy import moveaxis, _promote_args_inexact, _promote_dtypes_inexact
 from jax._src.numpy.util import _wraps
 from jax._src.ops import special as ops_special
+from jax._src.third_party.scipy.betaln import betaln as _betaln_impl
 from jax._src.typing import Array, ArrayLike
 
 
@@ -38,10 +39,11 @@ def gammaln(x: ArrayLike) -> Array:
   return lax.lgamma(x)
 
 
-@_wraps(osp_special.betaln, module='scipy.special')
-def betaln(x: ArrayLike, y: ArrayLike) -> Array:
-  x, y = _promote_args_inexact("betaln", x, y)
-  return lax.lgamma(x) + lax.lgamma(y) - lax.lgamma(x + y)
+betaln = _wraps(
+    osp_special.betaln,
+    module='scipy.special',
+    update_doc=False
+)(_betaln_impl)
 
 
 @_wraps(osp_special.betainc, module='scipy.special')
diff --git a/jax/_src/third_party/scipy/betaln.py b/jax/_src/third_party/scipy/betaln.py
new file mode 100644
--- /dev/null
+++ b/jax/_src/third_party/scipy/betaln.py
@@ -0,0 +1,62 @@
+from jax import lax
+import jax.numpy as jnp
+from jax._src.typing import Array, ArrayLike
+from jax._src.numpy.lax_numpy import _promote_args_inexact
+
+def algdiv(a: ArrayLike, b: ArrayLike) -> Array:
+    """
+    Compute ``log(gamma(a))/log(gamma(a + b))`` when ``b >= 8``.
+
+    Derived from scipy's implmentation of `algdiv`_.
+
+    This differs from the scipy implementation in that it assumes a <= b
+    because recomputing ``a, b = jnp.minimum(a, b), jnp.maximum(a, b)`` might
+    be expensive and this is only called by ``betaln``.
+
+    .. _algdiv:
+        https://github.com/scipy/scipy/blob/c89dfc2b90d993f2a8174e57e0cbc8fbe6f3ee19/scipy/special/cdflib/algdiv.f
+    """
+    c0 = 0.833333333333333e-01
+    c1 = -0.277777777760991e-02
+    c2 = 0.793650666825390e-03
+    c3 = -0.595202931351870e-03
+    c4 = 0.837308034031215e-03
+    c5 = -0.165322962780713e-02
+    h = a / b
+    c = h / (1 + h)
+    x = h / (1 + h)
+    d = b + (a - 0.5)
+    # Set sN = (1 - x**n)/(1 - x)
+    x2 = x * x
+    s3 = 1.0 + (x + x2)
+    s5 = 1.0 + (x + x2 * s3)
+    s7 = 1.0 + (x + x2 * s5)
+    s9 = 1.0 + (x + x2 * s7)
+    s11 = 1.0 + (x + x2 * s9)
+    # Set w = del(b) - del(a + b)
+    # where del(x) is defined by ln(gamma(x)) = (x - 0.5)*ln(x) - x + 0.5*ln(2*pi) + del(x)
+    t = (1.0 / b) ** 2
+    w = ((((c5 * s11 * t + c4 * s9) * t + c3 * s7) * t + c2 * s5) * t + c1 * s3) * t + c0
+    w = w * (c / b)
+    # Combine the results
+    u = d * lax.log1p(a / b)
+    v = a * (lax.log(b) - 1.0)
+    return jnp.where(u <= v, (w - v) - u, (w - u) - v)
+
+
+def betaln(a: ArrayLike, b: ArrayLike) -> Array:
+    """Compute the log of the beta function.
+
+    Derived from scipy's implementation of `betaln`_.
+
+    This implementation does not handle all branches of the scipy implementation, but is still much more accurate
+    than just doing lgamma(a) + lgamma(b) - lgamma(a + b) when inputs are large (> 1M or so).
+
+    .. _betaln:
+        https://github.com/scipy/scipy/blob/ef2dee592ba8fb900ff2308b9d1c79e4d6a0ad8b/scipy/special/cdflib/betaln.f
+    """
+    a, b = _promote_args_inexact("betaln", a, b)
+    a, b = jnp.minimum(a, b), jnp.maximum(a, b)
+    small_b = lax.lgamma(a) + (lax.lgamma(b) - lax.lgamma(a + b))
+    large_b = lax.lgamma(a) + algdiv(a, b)
+    return jnp.where(b < 8, small_b, large_b)


Test patch:
diff --git a/tests/lax_scipy_test.py b/tests/lax_scipy_test.py
--- a/tests/lax_scipy_test.py
+++ b/tests/lax_scipy_test.py
@@ -303,6 +303,18 @@ def testIssue3758(self):
     q = np.array([1., 40., 30.], dtype=np.float32)
     self.assertAllClose(np.array([1., 0., 0.], dtype=np.float32), lsp_special.zeta(x, q))
 
+  def testIssue13267(self):
+    """Tests betaln(x, 1) across wide range of x."""
+    xs = jnp.geomspace(1, 1e30, 1000)
+    primals_out, tangents_out = jax.jvp(lsp_special.betaln, primals=[xs, 1.0], tangents=[jnp.ones_like(xs), 0.0])
+    # Check that betaln(x, 1) = -log(x).
+    # Betaln is still not perfect for small values, hence the atol (but it's close)
+    atol = jtu.if_device_under_test("tpu", 1e-3, 1e-5)
+    self.assertAllClose(primals_out, -jnp.log(xs), atol=atol)
+    # Check that d/dx betaln(x, 1) = d/dx -log(x) = -1/x.
+    self.assertAllClose(tangents_out, -1 / xs, atol=atol)
+
+
   def testXlogyShouldReturnZero(self):
     self.assertAllClose(lsp_special.xlogy(0., 0.), 0., check_dtypes=False)