jax.scipy.stats.chi2.logpdf returns NaN at x=0 and x=inf when df=2 (inconsistent with gamma.logpdf)
### Description
`jax.scipy.stats.chi2.logpdf(x, df)` returns `nan` where `scipy.stats.chi2.logpdf` returns a finite value (or `-inf`), whenever `(df/2 - 1) * log(y)` evaluates to `0 * (±inf)`:
```python
import jax
import jax.scipy.stats
import scipy.stats
jax.config.update('jax_enable_x64', True)
print(jax.scipy.stats.chi2.logpdf(0.0, 2.0)) # nan
print(scipy.stats.chi2.logpdf(0.0, 2.0)) # -0.6931471805599453
print(jax.scipy.stats.chi2.pdf(0.0, 2.0)) # nan
print(scipy.stats.chi2.pdf(0.0, 2.0)) # 0.5
print(jax.scipy.stats.chi2.logpdf(jax.numpy.inf, 2.0)) # nan (scipy: -inf)
```
The chi-square density with `df=2` is the exponential density `0.5 * exp(-x/2)`, which is well-defined at `x=0` — this is not a singularity of the distribution. `torch.distributions.Chi2(2.0).log_prob(0.0)` also returns `-0.6931...`, agreeing with scipy.
### Root cause
`jax/_src/scipy/stats/chi2.py` computes the kernel with a plain multiply:
```python
kernel = lax.sub(lax.mul(lax.sub(df_on_two, one), lax.log(y)), lax.div(y,two))
```
At `df=2`, `df/2 - 1 == 0`, so `x=0` gives `0 * (-inf) = nan` and `x=inf` gives `0 * inf = nan` (df>2 also hits `inf - inf = nan` at `x=inf`).
Notably, `jax/_src/scipy/stats/gamma.py::logpdf` already handles this exact case with `xlogy`:
```python
log_linear_term = lax.sub(xlogy(lax.sub(a, one), y), y)
```
so `jax.scipy.stats.gamma.logpdf(0.0, a=1)` correctly returns `0.0` while the mathematically identical `chi2.logpdf(0.0, df=2)` returns `nan`.
**Suggested fix:** use `xlogy(df/2 - 1, y)` in `chi2.logpdf`, mirroring `gamma.logpdf`. Reproduces in float32 and float64, on 0.10.1 and current `main` (the kernel line is unchanged at HEAD).
### System info (python version, jaxlib version, accelerator, etc.)
```
jax: 0.10.1
jaxlib: 0.10.1
numpy: 2.4.6
python: 3.12.13
device: CPU (arm64, macOS)
```
1 条评论