Merge pull request #18683 from jakevdp:gamma-neg

PiperOrigin-RevId: 585764804

jax/_src/scipy/special.py

@@ -46,8 +46,10 @@ def gammaln(x: ArrayLike) -> Array:
 The JAX version only accepts real-valued inputs.""")
 def gamma(x: ArrayLike) -> Array:
   x, = promote_args_inexact("gamma", x)
-  return lax.exp(lax.lgamma(x))
-
+  # Compute the sign for negative x, matching the semantics of scipy.special.gamma
+  floor_x = lax.floor(x)
+  sign = jnp.where((x > 0) | (x == floor_x), 1.0, (-1.0) ** floor_x)
+  return sign * lax.exp(lax.lgamma(x))
 
 betaln = _wraps(
     osp_special.betaln,

tests/lax_scipy_special_functions_test.py

@@ -62,7 +62,7 @@ def op_record(name, nargs, dtypes, rng_factory, test_grad, nondiff_argnums=(), t
         "betainc", 3, float_dtypes, jtu.rand_positive, False
     ),
     op_record(
-        "gamma", 1, float_dtypes, jtu.rand_positive, True
+        "gamma", 1, float_dtypes, jtu.rand_default, True
     ),
     op_record(
         "digamma", 1, float_dtypes, jtu.rand_positive, True
@@ -199,6 +199,15 @@ def testScipySpecialFunBernoulli(self, n):
     self._CheckAgainstNumpy(scipy_op, lax_op, args_maker, atol=0, rtol=1E-5)
     self._CompileAndCheck(lax_op, args_maker, atol=0, rtol=1E-5)
 
+  def testGammaSign(self):
+    # Test that the sign of `gamma` matches at integer-valued inputs.
+    dtype = jax.numpy.zeros(0).dtype  # default float dtype.
+    args_maker = lambda: [np.arange(-10, 10).astype(dtype)]
+    rtol = 1E-3 if jtu.test_device_matches(["tpu"]) else 1e-5
+    self._CheckAgainstNumpy(osp_special.gamma, lsp_special.gamma, args_maker, rtol=rtol)
+    self._CompileAndCheck(lsp_special.gamma, args_maker, rtol=rtol)
+
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())