Repository: google/jax

Issue:
Incorrect gradient for `jax.scipy.stats.beta.logpdf` 
### Description

Hi all. During some testing for an inference problem, I noticed that the gradients computed for `jax.scipy.stats.beta.logpdf` were incorrect when using alpha/beta parameters at 1. Please find a MWE below.

```python
import jax
import jax.numpy as jnp
import jax.scipy.special as jaxspec
import jax.scipy.stats as jaxstats

from jax import Array
from jax.config import config
from jax.scipy.special import gammaln
from jax.typing import ArrayLike
config.update("jax_enable_x64", True)


# custom loglikelihood
def jax_loglik(params: ArrayLike, p: ArrayLike) -> Array:
    k, n = params
    return (
        (k - 1) * jnp.sum(jnp.log(p))
        + (n - 1) * jnp.sum(jnp.log1p(-p))
        - len(p) * (gammaln(k) + gammaln(n) - gammaln(k + n))
    )

# log likelihood using direct logpdf sum
def jax_logpdf(params: ArrayLike, p: ArrayLike) -> Array:
    k, n = params
    return jnp.sum(jaxstats.beta.logpdf(p, k, n))

# fake data
x = 0.5 * jnp.ones(2)

# problematic parameters
bad_params = jnp.ones(2)

# non-problematic parameters
fine_params = 2 * jnp.ones(2)

for name, param in [("bad param", bad_params), ("fine param", fine_params)]:
    print(f"Evaluating functions using {name}")

    # eval each log likelihood definition using alpha = 1., beta = 1., x1 = 0.5, x2 = 0.5
    print(f"Bespoke ll = {jax_loglik(param, x)}")
    print(f"JAX beta.logpdf = {jax_logpdf(param, x)}")
    
    # eval gradient of log likelihood wrt alpha, beta, at alpha = beta = 1.
    print(f"Grad bespoke ll = {jax.grad(jax_loglik)(param, x)}")
    print(f"Grad JAX beta.logpdf {jax.grad(jax_logpdf)(param, x)}")
    print("")
```

The above code generates the following output, 
```
Evaluating functions using bad param
Bespoke ll = -0.0
JAX beta.logpdf = 0.0
Grad bespoke ll = [0.61370564 0.61370564]
Grad JAX beta.logpdf [2. 2.]

Evaluating functions using fine param
Bespoke ll = 0.8109302162163301
JAX beta.logpdf = 0.8109302162163301
Grad bespoke ll = [0.28037231 0.28037231]
Grad JAX beta.logpdf [0.28037231 0.28037231]
```

### What jax/jaxlib version are you using?

jax v0.4.8

### Which accelerator(s) are you using?

CPU

### Additional system info

Mac M1

### NVIDIA GPU info

_No response_


Discussion:
I've edited the above example to be a bit more minimal and include a positive example where gradients agree for non-problematic parameters.
After a bit of digging, it looks like this is related to the use of `xlogy` and `xlog1py` in the definition of `logpdf`. Here is a reduced MWE that pertains to those parts:

```python
import jax
import jax.numpy as jnp
import jax.scipy.stats as jaxstats

from jax import Array
from jax.config import config
from jax.scipy.special import betaln, gammaln, polygamma, xlogy, xlog1py
from jax.typing import ArrayLike

config.update("jax_enable_x64", True)


def jax_norm(params: ArrayLike, p: ArrayLike) -> Array:
    k, n = params
    return -len(p) * betaln(k, n)


def custom_norm_grad(params: ArrayLike, p: ArrayLike) -> Array:
    k, n = params
    return jnp.array(
        [
            -len(p) * polygamma(0, k) + len(p) * polygamma(0, k + n),
            -len(p) * polygamma(0, n) + len(p) * polygamma(0, k + n),
        ]
    )


def custom_unscaled_loglike_grad(params: ArrayLike, p: ArrayLike) -> Array:
    k, n = params
    return jnp.array([jnp.sum(jnp.log(p)), jnp.sum(jnp.log1p(-p))])


def jax_unscaled_loglike(params: ArrayLike, p: ArrayLike) -> Array:
    k, n = params
    return jnp.sum(xlogy((k - 1), jnp.log(p)) + xlog1py((n - 1), -p))


# fake data
x = 0.5 * jnp.ones(2)

# problematic parameters
param = jnp.ones(2)

print(
    f"Bespoke gradient of unscaled log-likelihood {custom_unscaled_loglike_grad(param, x)}"
)
print(
    f"JAX Gradient of unscaled log-likelihood {jax.grad(jax_unscaled_loglike)(param, x)}"
)

print(f"Bespoke gradient of log-likelihood normalizing term {custom_norm_grad(param, x)}")
print(f"JAX Gradient of log-likelihood normalizing term {jax.grad(jax_norm)(param, x)}")

```

which outputs,
```
Bespoke gradient of unscaled log-likelihood [-1.38629436 -1.38629436]
JAX Gradient of unscaled log-likelihood [0. 0.]
Bespoke gradient of log-likelihood normalizing term [2. 2.]
JAX Gradient of log-likelihood normalizing term [2. 2.]
```

So something is happening (maybe on forward pass?) that first evaluates the `xlogy` (or `xlog1py`) terms to 0, since `k - 1` and `n - 1` are 0 when `k = n = 1`, but the gradients should still evaluate to non-zero using `xlogy` (or `xlog1py`) with `x=1`.
Thanks - I think you're right in your assessment that `xlogy` and `xlog1py` have incorrect gradients at x=0. An easier way to demonstrate this:
```python
print(jax.grad(lambda x: xlogy(x, 2.0))(0.0))
# 0.0
print(jax.grad(lambda x: x * jnp.log(2.0))(0.0))
# 0.6931472
```
These probably need a custom JVP given their current implementation.

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
@@ -114,23 +114,40 @@ def expit(x: ArrayLike) -> Array:
 logsumexp = _wraps(osp_special.logsumexp, module='scipy.special')(ops_special.logsumexp)
 
 
+@custom_derivatives.custom_jvp
 @_wraps(osp_special.xlogy, module='scipy.special')
 def xlogy(x: ArrayLike, y: ArrayLike) -> Array:
+  # Note: xlogy(0, 0) should return 0 according to the function documentation.
   x, y = promote_args_inexact("xlogy", x, y)
   x_ok = x != 0.
   safe_x = jnp.where(x_ok, x, 1.)
   safe_y = jnp.where(x_ok, y, 1.)
   return jnp.where(x_ok, lax.mul(safe_x, lax.log(safe_y)), jnp.zeros_like(x))
 
+def _xlogy_jvp(primals, tangents):
+  (x, y) = primals
+  (x_dot, y_dot) = tangents
+  result = xlogy(x, y)
+  return result, (x_dot * lax.log(y) + y_dot * x / y).astype(result.dtype)
+xlogy.defjvp(_xlogy_jvp)
+
 
+@custom_derivatives.custom_jvp
 @_wraps(osp_special.xlog1py, module='scipy.special', update_doc=False)
 def xlog1py(x: ArrayLike, y: ArrayLike) -> Array:
+  # Note: xlog1py(0, -1) should return 0 according to the function documentation.
   x, y = promote_args_inexact("xlog1py", x, y)
   x_ok = x != 0.
   safe_x = jnp.where(x_ok, x, 1.)
   safe_y = jnp.where(x_ok, y, 1.)
   return jnp.where(x_ok, lax.mul(safe_x, lax.log1p(safe_y)), jnp.zeros_like(x))
 
+def _xlog1py_jvp(primals, tangents):
+  (x, y) = primals
+  (x_dot, y_dot) = tangents
+  result = xlog1py(x, y)
+  return result, (x_dot * lax.log1p(y) + y_dot * x / (1 + y)).astype(result.dtype)
+xlog1py.defjvp(_xlog1py_jvp)
 
 @_wraps(osp_special.entr, module='scipy.special')
 def entr(x: ArrayLike) -> Array:


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
@@ -13,7 +13,6 @@
 # limitations under the License.
 
 
-import functools
 from functools import partial
 import itertools
 import unittest
@@ -220,15 +219,29 @@ def testXlogyShouldReturnZero(self):
     self.assertAllClose(lsp_special.xlogy(0., 0.), 0., check_dtypes=False)
 
   def testGradOfXlogyAtZero(self):
-    partial_xlogy = functools.partial(lsp_special.xlogy, 0.)
-    self.assertAllClose(jax.grad(partial_xlogy)(0.), 0., check_dtypes=False)
+    # https://github.com/google/jax/issues/15598
+    x0, y0 = 0.0, 3.0
+    d_xlog1py_dx = jax.grad(lsp_special.xlogy, argnums=0)(x0, y0)
+    self.assertAllClose(d_xlog1py_dx, lax.log(y0))
+
+    d_xlog1py_dy = jax.grad(lsp_special.xlogy, argnums=1)(x0, y0)
+    self.assertAllClose(d_xlog1py_dy, 0.0)
+
+    jtu.check_grads(lsp_special.xlogy, (x0, y0), order=2)
 
   def testXlog1pyShouldReturnZero(self):
     self.assertAllClose(lsp_special.xlog1py(0., -1.), 0., check_dtypes=False)
 
   def testGradOfXlog1pyAtZero(self):
-    partial_xlog1py = functools.partial(lsp_special.xlog1py, 0.)
-    self.assertAllClose(jax.grad(partial_xlog1py)(-1.), 0., check_dtypes=False)
+    # https://github.com/google/jax/issues/15598
+    x0, y0 = 0.0, 3.0
+    d_xlog1py_dx = jax.grad(lsp_special.xlog1py, argnums=0)(x0, y0)
+    self.assertAllClose(d_xlog1py_dx, lax.log1p(y0))
+
+    d_xlog1py_dy = jax.grad(lsp_special.xlog1py, argnums=1)(x0, y0)
+    self.assertAllClose(d_xlog1py_dy, 0.0)
+
+    jtu.check_grads(lsp_special.xlog1py, (x0, y0), order=2)
 
   @jtu.sample_product(
     [dict(order=order, z=z, n_iter=n_iter)