Repository: google/jax

Issue:
Numerically inaccurate `scipy.stats.norm.sf`
### Description

Since a distribution cdf's goes asymptotically to 1 as $x \to \infty$, computing it numerically rounds to 1 at a certain point. For this reason, distributions in `scipy.stats` also provide the method `sf = 1 - cdf` that retains the precision.

I was about to use `jax.scipy.stats.norm.sf`, but I noticed that it is not accurate:

```python
from jax.scipy import stats as jstats
from scipy import stats
import numpy as np

value = np.array(10, np.float32)
print('scipy:', stats.norm.sf(value))
print('jax:', jstats.norm.sf(value))
```

```
scipy: 7.61985302416047e-24
jax: 0.0
```

Indeed, looking at the code, the implementation just does `1 - cdf`:

```python
@_wraps(osp_stats.norm.sf, update_doc=False)
def sf(x: ArrayLike, loc: ArrayLike = 0, scale: ArrayLike = 1) -> Array:
  cdf_result = cdf(x, loc, scale)
  return lax.sub(_lax_const(cdf_result, 1), cdf_result)
```

For the normal distribution, one can use `sf(x) = cdf(-x)`. I have not looked at what jax is doing for other distributions.

### What jax/jaxlib version are you using?

0.4.14

### Which accelerator(s) are you using?

CPU

### Additional system info

Python 3.11.2, macOS 13.4

### NVIDIA GPU info

_No response_


Discussion:
Thanks for the report! #17204 fixes this for `norm.sf`; I'll look into similar fixes for the other distributions as well.

Resolving patch:
diff --git a/jax/_src/scipy/stats/norm.py b/jax/_src/scipy/stats/norm.py
--- a/jax/_src/scipy/stats/norm.py
+++ b/jax/_src/scipy/stats/norm.py
@@ -57,10 +57,16 @@ def ppf(q: ArrayLike, loc: ArrayLike = 0, scale: ArrayLike = 1) -> Array:
   return jnp.asarray(special.ndtri(q) * scale + loc, float)
 
 
+@_wraps(osp_stats.norm.logsf, update_doc=False)
+def logsf(x: ArrayLike, loc: ArrayLike = 0, scale: ArrayLike = 1) -> Array:
+  x, loc, scale = promote_args_inexact("norm.logsf", x, loc, scale)
+  return logcdf(-x, -loc, scale)
+
+
 @_wraps(osp_stats.norm.sf, update_doc=False)
 def sf(x: ArrayLike, loc: ArrayLike = 0, scale: ArrayLike = 1) -> Array:
-  cdf_result = cdf(x, loc, scale)
-  return lax.sub(_lax_const(cdf_result, 1), cdf_result)
+  x, loc, scale = promote_args_inexact("norm.sf", x, loc, scale)
+  return cdf(-x, -loc, scale)
 
 
 @_wraps(osp_stats.norm.isf, update_doc=False)
diff --git a/jax/scipy/stats/norm.py b/jax/scipy/stats/norm.py
--- a/jax/scipy/stats/norm.py
+++ b/jax/scipy/stats/norm.py
@@ -19,6 +19,7 @@
   cdf as cdf,
   logcdf as logcdf,
   logpdf as logpdf,
+  logsf as logsf,
   pdf as pdf,
   ppf as ppf,
   sf as sf,


Test patch:
diff --git a/tests/scipy_stats_test.py b/tests/scipy_stats_test.py
--- a/tests/scipy_stats_test.py
+++ b/tests/scipy_stats_test.py
@@ -693,6 +693,23 @@ def args_maker():
                               tol=1e-6)
       self._CompileAndCheck(lax_fun, args_maker)
 
+  @genNamedParametersNArgs(3)
+  def testNormLogSf(self, shapes, dtypes):
+    rng = jtu.rand_default(self.rng())
+    scipy_fun = osp_stats.norm.logsf
+    lax_fun = lsp_stats.norm.logsf
+
+    def args_maker():
+      x, loc, scale = map(rng, shapes, dtypes)
+      # clipping to ensure that scale is not too low
+      scale = np.clip(np.abs(scale), a_min=0.1, a_max=None).astype(scale.dtype)
+      return [x, loc, scale]
+
+    with jtu.strict_promotion_if_dtypes_match(dtypes):
+      self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker, check_dtypes=False,
+                              tol=1e-4)
+      self._CompileAndCheck(lax_fun, args_maker)
+
   @genNamedParametersNArgs(3)
   def testNormSf(self, shapes, dtypes):
     rng = jtu.rand_default(self.rng())
@@ -710,6 +727,13 @@ def args_maker():
                               tol=1e-6)
       self._CompileAndCheck(lax_fun, args_maker)
 
+  def testNormSfNearZero(self):
+    # Regression test for https://github.com/google/jax/issues/17199
+    value = np.array(10, np.float32)
+    self.assertAllClose(osp_stats.norm.sf(value).astype('float32'),
+                        lsp_stats.norm.sf(value),
+                        atol=0, rtol=1E-5)
+
   @genNamedParametersNArgs(3)
   def testNormPpf(self, shapes, dtypes):
     rng = jtu.rand_default(self.rng())