prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  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
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `jax/_src/scipy/stats/binom.py`
  Content:
  ```
  1 # Copyright 2023 The JAX Authors.
  2 #
  3 # Licensed under the Apache License, Version 2.0 (the "License");
  4 # you may not use this file except in compliance with the License.
  5 # You may obtain a copy of the License at
  6 #
  7 #     https://www.apache.org/licenses/LICENSE-2.0
  8 #
  9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License
  14 
  15 
  16 import scipy.stats as osp_stats
  17 
  18 from jax import lax
  19 import jax.numpy as jnp
  20 from jax._src.numpy.util import _wraps, promote_args_inexact
  21 from jax._src.scipy.special import gammaln, xlogy, xlog1py
  22 from jax._src.typing import Array, ArrayLike
  23 
  24 
  25 @_wraps(osp_stats.nbinom.logpmf, update_doc=False)
  26 def logpmf(k: ArrayLike, n: ArrayLike, p: ArrayLike, loc: ArrayLike = 0) -> Array:
  27     """JAX implementation of scipy.stats.binom.logpmf."""
  28     k, n, p, loc = promote_args_inexact("binom.logpmf", k, n, p, loc)
  29     y = lax.sub(k, loc)
  30     comb_term = lax.sub(
  31         gammaln(n + 1),
  32         lax.add(gammaln(y + 1), gammaln(n - y + 1))
  33     )
  34     log_linear_term = lax.add(xlogy(y, p), xlog1py(lax.sub(n, y), lax.neg(p)))
  35     log_probs = lax.add(comb_term, log_linear_term)
  36     return jnp.where(lax.lt(k, loc), -jnp.inf, log_probs)
  37 
  38 
  39 @_wraps(osp_stats.nbinom.pmf, update_doc=False)
  40 def pmf(k: ArrayLike, n: ArrayLike, p: ArrayLike, loc: ArrayLike = 0) -> Array:
  41     """JAX implementation of scipy.stats.binom.pmf."""
  42     return lax.exp(logpmf(k, n, p, loc))
  43
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/jax/_src/scipy/stats/binom.py b/jax/_src/scipy/stats/binom.py
  --- a/jax/_src/scipy/stats/binom.py
  +++ b/jax/_src/scipy/stats/binom.py
  @@ -33,7 +33,7 @@
       )
       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)
  
verification_info: {"golden_diff": "diff --git a/jax/_src/scipy/stats/binom.py b/jax/_src/scipy/stats/binom.py\n--- a/jax/_src/scipy/stats/binom.py\n+++ b/jax/_src/scipy/stats/binom.py\n@@ -33,7 +33,7 @@\n     )\n     log_linear_term = lax.add(xlogy(y, p), xlog1py(lax.sub(n, y), lax.neg(p)))\n     log_probs = lax.add(comb_term, log_linear_term)\n-    return jnp.where(lax.lt(k, loc), -jnp.inf, log_probs)\n+    return jnp.where(lax.ge(k, loc) & lax.lt(k, loc + n + 1), log_probs, -jnp.inf)\n \n \n @_wraps(osp_stats.nbinom.pmf, update_doc=False)\n", "issue": "Unexpected behavior of `jax.scipy.stats.binom.pmf`\n### Description\r\n\r\npmf 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.\r\n\r\n```python\r\nimport jax\r\nfrom jax import numpy as jnp\r\nfrom matplotlib import pyplot as plt\r\n\r\nx = jnp.linspace(-1, 10, 1000)\r\nxxf = jax.scipy.stats.binom.pmf(k=x, n=5, p=0.8)\r\n\r\nplt.plot(x, xxf)\r\nplt.tight_layout()\r\nplt.show()\r\n```\r\n![xyxy](https://github.com/google/jax/assets/62182585/c35a2ebb-e914-4b37-9f59-6ccd7cc10b28)\r\nThe side left to the zero is as expected.\r\n\r\n### What jax/jaxlib version are you using?\r\n\r\njax v0.4.23\r\n\r\n### Which accelerator(s) are you using?\r\n\r\nCPU\n", "before_files": [{"content": "# Copyright 2023 The JAX Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License\n\n\nimport scipy.stats as osp_stats\n\nfrom jax import lax\nimport jax.numpy as jnp\nfrom jax._src.numpy.util import _wraps, promote_args_inexact\nfrom jax._src.scipy.special import gammaln, xlogy, xlog1py\nfrom jax._src.typing import Array, ArrayLike\n\n\n@_wraps(osp_stats.nbinom.logpmf, update_doc=False)\ndef logpmf(k: ArrayLike, n: ArrayLike, p: ArrayLike, loc: ArrayLike = 0) -> Array:\n    \"\"\"JAX implementation of scipy.stats.binom.logpmf.\"\"\"\n    k, n, p, loc = promote_args_inexact(\"binom.logpmf\", k, n, p, loc)\n    y = lax.sub(k, loc)\n    comb_term = lax.sub(\n        gammaln(n + 1),\n        lax.add(gammaln(y + 1), gammaln(n - y + 1))\n    )\n    log_linear_term = lax.add(xlogy(y, p), xlog1py(lax.sub(n, y), lax.neg(p)))\n    log_probs = lax.add(comb_term, log_linear_term)\n    return jnp.where(lax.lt(k, loc), -jnp.inf, log_probs)\n\n\n@_wraps(osp_stats.nbinom.pmf, update_doc=False)\ndef pmf(k: ArrayLike, n: ArrayLike, p: ArrayLike, loc: ArrayLike = 0) -> Array:\n    \"\"\"JAX implementation of scipy.stats.binom.pmf.\"\"\"\n    return lax.exp(logpmf(k, n, p, loc))\n", "path": "jax/_src/scipy/stats/binom.py"}], "after_files": [{"content": "# Copyright 2023 The JAX Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License\n\n\nimport scipy.stats as osp_stats\n\nfrom jax import lax\nimport jax.numpy as jnp\nfrom jax._src.numpy.util import _wraps, promote_args_inexact\nfrom jax._src.scipy.special import gammaln, xlogy, xlog1py\nfrom jax._src.typing import Array, ArrayLike\n\n\n@_wraps(osp_stats.nbinom.logpmf, update_doc=False)\ndef logpmf(k: ArrayLike, n: ArrayLike, p: ArrayLike, loc: ArrayLike = 0) -> Array:\n    \"\"\"JAX implementation of scipy.stats.binom.logpmf.\"\"\"\n    k, n, p, loc = promote_args_inexact(\"binom.logpmf\", k, n, p, loc)\n    y = lax.sub(k, loc)\n    comb_term = lax.sub(\n        gammaln(n + 1),\n        lax.add(gammaln(y + 1), gammaln(n - y + 1))\n    )\n    log_linear_term = lax.add(xlogy(y, p), xlog1py(lax.sub(n, y), lax.neg(p)))\n    log_probs = lax.add(comb_term, log_linear_term)\n    return jnp.where(lax.ge(k, loc) & lax.lt(k, loc + n + 1), log_probs, -jnp.inf)\n\n\n@_wraps(osp_stats.nbinom.pmf, update_doc=False)\ndef pmf(k: ArrayLike, n: ArrayLike, p: ArrayLike, loc: ArrayLike = 0) -> Array:\n    \"\"\"JAX implementation of scipy.stats.binom.pmf.\"\"\"\n    return lax.exp(logpmf(k, n, p, loc))\n", "path": "jax/_src/scipy/stats/binom.py"}]}