# Copyright 2023 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License


import scipy.stats as osp_stats

from jax import lax
import jax.numpy as jnp
from jax._src.numpy.util import _wraps, promote_args_inexact
from jax._src.scipy.special import gammaln, xlogy, xlog1py
from jax._src.typing import Array, ArrayLike


@_wraps(osp_stats.nbinom.logpmf, update_doc=False)
def logpmf(k: ArrayLike, n: ArrayLike, p: ArrayLike, loc: ArrayLike = 0) -> Array:
    """JAX implementation of scipy.stats.binom.logpmf."""
    k, n, p, loc = promote_args_inexact("binom.logpmf", k, n, p, loc)
    y = lax.sub(k, loc)
    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)
    return jnp.where(lax.lt(k, loc), -jnp.inf, log_probs)


@_wraps(osp_stats.nbinom.pmf, update_doc=False)
def pmf(k: ArrayLike, n: ArrayLike, p: ArrayLike, loc: ArrayLike = 0) -> Array:
    """JAX implementation of scipy.stats.binom.pmf."""
    return lax.exp(logpmf(k, n, p, loc))

jax/_src/scipy/stats/binom.py
# Copyright 2023 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License


import scipy.stats as osp_stats

from jax import lax
import jax.numpy as jnp
from jax._src.numpy.util import _wraps, promote_args_inexact
from jax._src.scipy.special import gammaln, xlogy, xlog1py
from jax._src.typing import Array, ArrayLike


@_wraps(osp_stats.nbinom.logpmf, update_doc=False)
def logpmf(k: ArrayLike, n: ArrayLike, p: ArrayLike, loc: ArrayLike = 0) -> Array:
    """JAX implementation of scipy.stats.binom.logpmf."""
    k, n, p, loc = promote_args_inexact("binom.logpmf", k, n, p, loc)
    y = lax.sub(k, loc)
    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)
    return jnp.where(lax.ge(k, loc) & lax.lt(k, loc + n + 1), log_probs, -jnp.inf)


@_wraps(osp_stats.nbinom.pmf, update_doc=False)
def pmf(k: ArrayLike, n: ArrayLike, p: ArrayLike, loc: ArrayLike = 0) -> Array:
    """JAX implementation of scipy.stats.binom.pmf."""
    return lax.exp(logpmf(k, n, p, loc))

jax/_src/scipy/stats/binom.py
diff --git a/jax/_src/scipy/stats/binom.py b/jax/_src/scipy/stats/binom.py
index 852d39edc8e8..869eab91c5b5 100644
--- a/jax/_src/scipy/stats/binom.py
+++ b/jax/_src/scipy/stats/binom.py
@@ -33,7 +33,7 @@ def logpmf(k: ArrayLike, n: ArrayLike, p: ArrayLike, loc: ArrayLike = 0) -> Arra
     )
     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)
-    return jnp.where(lax.lt(k, loc), -jnp.inf, log_probs)
+    return jnp.where(lax.ge(k, loc) & lax.lt(k, loc + n + 1), log_probs, -jnp.inf)
 
 
 @_wraps(osp_stats.nbinom.pmf, update_doc=False)
diff --git a/tests/scipy_stats_test.py b/tests/scipy_stats_test.py
index 2db5b354a4cb..09b12b87b1bc 100644
--- a/tests/scipy_stats_test.py
+++ b/tests/scipy_stats_test.py
@@ -1159,6 +1159,10 @@ def args_maker():
                               tol=5e-4)
       self._CompileAndCheck(lax_fun, args_maker, rtol=tol, atol=tol)
 
+  def testBinomPmfOutOfRange(self):
+    # 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 testIssue972(self):
     self.assertAllClose(
       np.ones((4,), np.float32),

Unexpected behavior of `jax.scipy.stats.binom.pmf`
### Description

pmf of a random variable should be zero outside of its range. While plotting the graph for `jax.scipy.stats.binom.pmf`, I notice that for $n>5$ and $p>0.5$, there are some oscillations in the values of the pmf, which should not be there. For evidence, I am attaching a plot too.

```python
import jax
from jax import numpy as jnp
from matplotlib import pyplot as plt

x = jnp.linspace(-1, 10, 1000)
xxf = jax.scipy.stats.binom.pmf(k=x, n=5, p=0.8)

plt.plot(x, xxf)
plt.tight_layout()
plt.show()
```
![xyxy](https://github.com/google/jax/assets/62182585/c35a2ebb-e914-4b37-9f59-6ccd7cc10b28)
The side left to the zero is as expected.

### What jax/jaxlib version are you using?

jax v0.4.23

### Which accelerator(s) are you using?

CPU
