Repository: google/jax (python)

Issue:
Binomial and beta-binomial `logpmf` returns NaN and nonzero values when k=n=0 (it should be zero)
### Description

So I am building a model with `jax` + `jaxopt` and was getting some NaN's in my loss function (essentially the negative logpmf of the beta-binomial distribution) during training. Taking a look at the values going in, nothing seemed out of the ordinary, except I noticed that the NaN's coming out corresponded to scenarios where k=n=0. I'm pretty sure the output value of this should be zero, because the probability of getting zero successful events out of zero trials is 100% (and `log(1)==0`) :smile:.

I confirmed this with `scipy`'s version of the beta-binomial and found that it behaved as I expected. I then wondered whether this bug might also be present in the normal (non-beta) binomial and found that while the logpmf output is not NaN, it is generally a small float value rather than 0 exactly. Perhaps it's close enough to zero practically speaking that it is not an issue. You can reproduce the bug and see the comparison with `scipy` by running the following in an interpreter:

```python
Python 3.11.6 (main, Oct  3 2023, 02:51:45) [Clang 14.0.3 (clang-1403.0.22.14.1)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax.scipy.stats as jst
>>> import scipy.stats as sst
>>> import numpy as np
>>> # compare beta-binomial when k=n=0
>>> a = b = np.linspace(1, 100, 10)
>>> jst.betabinom.logpmf(0, 0, a, b)
Array([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan], dtype=float32)
>>> sst.betabinom.logpmf(0, 0, a, b)
array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
>>> # compare regular binomial when k=n=0
>>> p = np.linspace(0, 1, 10)
>>> jst.binom.logpmf(0, 0, p)
Array([-4.7683716e-07, -4.7683716e-07, -4.7683716e-07, -4.7683716e-07,
       -4.7683716e-07, -4.7683716e-07, -4.7683716e-07, -4.7683716e-07,
       -4.7683716e-07, -4.7683716e-07], dtype=float32)
>>> sst.binom.logpmf(0, 0, p)
array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
```

Looking at `jax`'s source, the problem with the beta-binomial seems to stem from [these lines](https://github.com/google/jax/blob/760deb310e7a5d574ceff03a0e38c807940b3f85/jax/_src/scipy/stats/betabinom.py#L39-L40), where the output is coerced to `jnp.nan` if `lax.lt(n, one)` is true. I think perhaps if you just replace `one` with `zero` here that should fix it. If you look at the [scipy implementation](https://github.com/scipy/scipy/blob/ae3ca70a72ef04547fbc28501009246cef5ee6c8/scipy/stats/_discrete_distns.py#L230-L233) that check isn't present at all. I tested it out locally and it still doesn't evaluate to zero exactly like the regular binomial, but I don't think it's that big of a deal.

```python
>>> logpmf(0, 0, a, b) # after replacing `one` with `zero` in `lax.lt(n, one)`
Array([-9.536743e-07, -9.536743e-07, -9.536743e-07, -9.536743e-07,
       -9.536743e-07, -9.536743e-07, -9.536743e-07, -9.536743e-07,
       -9.536743e-07, -9.536743e-07], dtype=float32)
```

I am happy to whip up a PR real quick if this makes sense!

### What jax/jaxlib version are you using?

jax==0.4.23, jaxlib==0.4.23

### Which accelerator(s) are you using?

CPU

### Additional system info?

Python 3.11.6, MacOS

### NVIDIA GPU info

_No response_

Pull request description:
Fix `jax.scipy.stats.betabinom.logpmf` and `binom.logpmf` for JAX to emulate SciPy's behavior when `k=n=0`
Fixes #19318 

Patch:
diff --git a/jax/_src/scipy/stats/betabinom.py b/jax/_src/scipy/stats/betabinom.py
index d526c373f..7d4a4ed79 100644
--- a/jax/_src/scipy/stats/betabinom.py
+++ b/jax/_src/scipy/stats/betabinom.py
@@ -55,9 +55,11 @@ def logpmf(k: ArrayLike, n: ArrayLike, a: ArrayLike, b: ArrayLike,
   combiln = lax.neg(lax.add(lax.log1p(n), betaln(lax.add(lax.sub(n,y), one), lax.add(y,one))))
   beta_lns = lax.sub(betaln(lax.add(y,a), lax.add(lax.sub(n,y),b)), betaln(a,b))
   log_probs = lax.add(combiln, beta_lns)
-  y_cond = jnp.logical_or(lax.lt(y, lax.neg(loc)), lax.gt(y, lax.sub(n, loc)))
+  log_probs = jnp.where(jnp.logical_and(lax.eq(y, zero), lax.eq(n, zero)), 0., log_probs)
+  y_cond = jnp.logical_or(jnp.logical_or(lax.lt(y, lax.neg(loc)), lax.gt(y, n)),
+                          lax.le(lax.add(y, a), zero))
   log_probs = jnp.where(y_cond, -jnp.inf, log_probs)
-  n_a_b_cond = jnp.logical_or(jnp.logical_or(lax.lt(n, one), lax.lt(a, zero)), lax.lt(b, zero))
+  n_a_b_cond = jnp.logical_or(jnp.logical_or(lax.lt(n, zero), lax.le(a, zero)), lax.le(b, zero))
   return jnp.where(n_a_b_cond, jnp.nan, log_probs)
 
 
diff --git a/jax/_src/scipy/stats/binom.py b/jax/_src/scipy/stats/binom.py
index e97e84d1c..50d5aa6a8 100644
--- a/jax/_src/scipy/stats/binom.py
+++ b/jax/_src/scipy/stats/binom.py
@@ -15,6 +15,7 @@
 from jax import lax
 import jax.numpy as jnp
 from jax._src.numpy.util import promote_args_inexact
+from jax._src.lax.lax import _const as _lax_const
 from jax._src.scipy.special import gammaln, xlogy, xlog1py
 from jax._src.typing import Array, ArrayLike
 
@@ -46,12 +47,16 @@ def logpmf(k: ArrayLike, n: ArrayLike, p: ArrayLike, loc: ArrayLike = 0) -> Arra
   """
   k, n, p, loc = promote_args_inexact("binom.logpmf", k, n, p, loc)
   y = lax.sub(k, loc)
+  zero = _lax_const(y, 0)
   comb_term = lax.sub(
       gammaln(n + 1),
       lax.add(gammaln(y + 1), gammaln(n - y + 1))
   )
   log_linear_term = lax.add(xlogy(y, p), xlog1py(lax.sub(n, y), lax.neg(p)))
   log_probs = lax.add(comb_term, log_linear_term)
+  y_n_cond = jnp.logical_or(jnp.logical_and(lax.eq(y, zero), lax.eq(n, zero)),
+                            lax.eq(log_linear_term, zero))
+  log_probs = jnp.where(y_n_cond, 0., log_probs)
   return jnp.where(lax.ge(k, loc) & lax.lt(k, loc + n + 1), log_probs, -jnp.inf)
 
 


Test patch:
diff --git a/tests/scipy_stats_test.py b/tests/scipy_stats_test.py
index 501f4cbe5..ad988bba6 100644
--- a/tests/scipy_stats_test.py
+++ b/tests/scipy_stats_test.py
@@ -1173,8 +1173,6 @@ class LaxBackedScipyStatsTests(jtu.JaxTestCase):
       k, n, a, b, loc = map(rng, shapes, dtypes)
       k = np.floor(k)
       n = np.ceil(n)
-      a = np.clip(a, a_min = 0.1, a_max=None).astype(a.dtype)
-      b = np.clip(a, a_min = 0.1, a_max=None).astype(b.dtype)
       loc = np.floor(loc)
       return [k, n, a, b, loc]
 
@@ -1184,6 +1182,10 @@ class LaxBackedScipyStatsTests(jtu.JaxTestCase):
                               tol=5e-4)
       self._CompileAndCheck(lax_fun, args_maker, rtol=1e-5, atol=1e-5)
 
+  def testBetaBinomLogPmfZerokZeron(self):
+    self.assertEqual(lsp_stats.betabinom.logpmf(0, 0, 10, 5, 0),
+                     osp_stats.betabinom.logpmf(0, 0, 10, 5, 0))
+
   @genNamedParametersNArgs(4)
   def testBinomLogPmf(self, shapes, dtypes):
     rng = jtu.rand_positive(self.rng())
@@ -1192,8 +1194,8 @@ class LaxBackedScipyStatsTests(jtu.JaxTestCase):
 
     def args_maker():
       k, n, logit, loc = map(rng, shapes, dtypes)
-      k = np.floor(np.abs(k))
-      n = np.ceil(np.abs(n))
+      k = np.floor(k)
+      n = np.ceil(n)
       p = expit(logit)
       loc = np.floor(loc)
       return [k, n, p, loc]
@@ -1209,6 +1211,10 @@ class LaxBackedScipyStatsTests(jtu.JaxTestCase):
     # Regression test for https://github.com/google/jax/issues/19150
     self.assertEqual(lsp_stats.binom.pmf(k=6.5, n=5, p=0.8), 0.0)
 
+  def testBinomLogPmfZerokZeron(self):
+    self.assertEqual(lsp_stats.binom.logpmf(0, 0, 0.8, 0),
+                     osp_stats.binom.logpmf(0, 0, 0.8, 0))
+
   def testIssue972(self):
     self.assertAllClose(
       np.ones((4,), np.float32),