Fix false positive debug_nans error caused by NaNs that are properly handled in jax.scipy.stats.gamma

As reported in jax-ml#24939, even though the implementation of `jax.scipy.stats.gamma.logpdf` handles invalid inputs (e.g. `x < loc`) by returning `-inf`, the existing implementation incorrectly triggers the NaN checks introduced by JAX's debug NaNs mode. This change updates the implementation to no longer produce internal NaNs.

Fixes jax-ml#24939

PiperOrigin-RevId: 698833589

Author: dfm

jax/_src/scipy/stats/gamma.py

@@ -51,12 +51,13 @@ def logpdf(x: ArrayLike, a: ArrayLike, loc: ArrayLike = 0, scale: ArrayLike = 1)
     - :func:`jax.scipy.stats.gamma.logsf`
   """
   x, a, loc, scale = promote_args_inexact("gamma.logpdf", x, a, loc, scale)
+  ok = lax.ge(x, loc)
   one = _lax_const(x, 1)
-  y = lax.div(lax.sub(x, loc), scale)
+  y = jnp.where(ok, lax.div(lax.sub(x, loc), scale), one)
   log_linear_term = lax.sub(xlogy(lax.sub(a, one), y), y)
   shape_terms = lax.add(gammaln(a), lax.log(scale))
   log_probs = lax.sub(log_linear_term, shape_terms)
-  return jnp.where(lax.lt(x, loc), -jnp.inf, log_probs)
+  return jnp.where(ok, log_probs, -jnp.inf)
 
 
 def pdf(x: ArrayLike, a: ArrayLike, loc: ArrayLike = 0, scale: ArrayLike = 1) -> Array:

tests/scipy_stats_test.py

@@ -543,6 +543,13 @@ def testGammaLogPdfZero(self):
     self.assertAllClose(
       osp_stats.gamma.pdf(0.0, 1.0), lsp_stats.gamma.pdf(0.0, 1.0), atol=1E-6)
 
+  def testGammaDebugNans(self):
+    # Regression test for https://github.com/jax-ml/jax/issues/24939
+    with jax.debug_nans(True):
+      self.assertAllClose(
+          osp_stats.gamma.pdf(0.0, 1.0, 1.0), lsp_stats.gamma.pdf(0.0, 1.0, 1.0)
+      )
+
   @genNamedParametersNArgs(4)
   def testGammaLogCdf(self, shapes, dtypes):
     rng = jtu.rand_positive(self.rng())