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_

Changed files:
tests/lax_scipy_test.py, jax/_src/scipy/special.py

diff --git a/jax/_src/scipy/special.py b/jax/_src/scipy/special.py
index f8f6d3f4b..1e732c021 100644
--- 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:
diff --git a/tests/lax_scipy_test.py b/tests/lax_scipy_test.py
index ca431eb9b..84d5da5ff 100644
--- 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 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
     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)