jax.scipy.special.kl_div returns nan instead of inf at x=inf, y=0.0

### Description `jax.scipy.special.kl_div(inf, 0.0)` returns `nan` instead of `inf`. The KL divergence is defined as: ``` kl_div(x, y) = x*log(x/y) - x + y if x > 0 and y > 0 = y if x == 0 and y >= 0 = inf otherwise ``` With `x=inf` and `y=0`, neither of the first two conditions holds, so the correct result is `inf`. SciPy returns `inf` correctly. JAX returns `nan`. ## Origin Found via testing derived from: ``` jax/tests/lax_scipy_special_functions_test.py ::LaxScipySpecialFunctionsTest::testScipySpecialFun_kl_div ``` The existing test uses `rand_positive` inputs only, which never include `x=inf` or `y=0`. The property uses a **reference oracle**: `jax.scipy.special.kl_div(x, y)` must agree with `scipy.special.kl_div(x, y)` in NaN/inf classification for all inputs. Related: #27088 (state of `jax.scipy.special` functions). ## Verification (SciPy is correct) The mathematical definition gives `+inf` for `x=inf, y=0`. The dominant term `x*log(x/y)` diverges, confirmed at finite approximations: ```python import numpy as np for x in [1e3, 1e6, 1e9, 1e12]: y = 1e-10 print(f"kl_div({x:.0e}, {y:.0e}) = {x * np.log(x/y) - x + y:.4e}") # kl_div(1e+03, 1e-10) = 2.8934e+04 # kl_div(1e+06, 1e-10) = 3.5841e+07 # kl_div(1e+09, 1e-10) = 4.2749e+10 # kl_div(1e+12, 1e-10) = 4.9657e+13 ``` The limit is `+∞`. SciPy is correct; the bug is in JAX. ## Minimal reproducer ```python import numpy as np import scipy.special import jax import jax.scipy.special jax.config.update("jax_enable_x64", True) x = np.float32(np.inf) y = np.float32(0.0) print(scipy.special.kl_div(x, y)) # inf ← correct print(jax.scipy.special.kl_div(x, y)) # nan ← BUG print(jax.jit(jax.scipy.special.kl_div)( jax.numpy.array(x), jax.numpy.array(y))) # nan ← BUG ``` ## Expected behavior ``` inf inf inf ``` ## Actual behavior ``` inf nan nan ``` ## Additional cases The bug is specific to `x=inf, y=0`. All other tested combinations are consistent with SciPy: ``` x=inf, y=0.0: scipy=inf, jax=nan ← BUG x=inf, y=0.5: scipy=nan, jax=nan OK x=inf, y=inf: scipy=nan, jax=nan OK x=0.0, y=0.0: scipy=0.0, jax=0.0 OK x=-1.0, y=0.5: scipy=inf, jax=inf OK x=nan, y=0.5: scipy=nan, jax=nan OK ``` ### System info (python version, jaxlib version, accelerator, etc.) ## System info ``` jax: 0.10.1 jaxlib: 0.10.1 numpy: 2.4.6 scipy: 1.17.1 python: 3.12.3 OS: Linux-6.17.0-29-generic-x86_64-with-glibc2.39 backend: cpu ```