Repository: google/jax

Issue:
`jax.scipy.special.ndtri` edge case behavior
### Description

`jax.scipy.special.ndtri` returns positive infinity when the argument exceeds `1.0` and negative infinity when the argument is less than `0.0`.

```python3
from jax import numpy as jnp
from jax.scipy.special import ndtri
ndtri(jnp.asarray(-1.))  # Array(-inf, dtype=float32)
ndtri(jnp.asarray(2.))  # Array(inf, dtype=float32)
```

It seems that it would be more appropriate to return NaN in these cases (as `scipy.special.ndtri`, `torch.special.ndtri`, etc. do). 

This came up when adding support for JAX as an array API backend in SciPy. For more information, please see scipy/scipy#20085.

### System info (python version, jaxlib version, accelerator, etc.)

```
jax:    0.4.23
jaxlib: 0.4.23
numpy:  1.25.2
python: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0]
jax.devices (1 total, 1 local): [cuda(id=0)]
process_count: 1

$ nvidia-smi
Tue Mar 26 04:33:34 2024       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.104.05             Driver Version: 535.104.05   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  Tesla T4                       Off | 00000000:00:04.0 Off |                    0 |
| N/A   70C    P0              32W /  70W |  11447MiB / 15360MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
+---------------------------------------------------------------------------------------+
```


Discussion:

Resolving patch:
diff --git a/jax/_src/scipy/special.py b/jax/_src/scipy/special.py
--- a/jax/_src/scipy/special.py
+++ b/jax/_src/scipy/special.py
@@ -540,7 +540,7 @@ def _create_polynomial(var, coeffs):
   # later on. The result from the computation when p == 0 is not used so any
   # number that doesn't result in NaNs is fine.
   sanitized_mcp = jnp.where(
-      maybe_complement_p <= dtype(0.),
+      maybe_complement_p == dtype(0.),
       jnp.full(shape, dtype(0.5)),
       maybe_complement_p)
 
@@ -571,9 +571,9 @@ def _create_polynomial(var, coeffs):
 
   x = jnp.where(p > dtype(1. - np.exp(-2.)), x, -x)
   infinity = jnp.full(shape, dtype(np.inf))
-  x_nan_replaced = jnp.where(
-      p <= dtype(0.0), -infinity, jnp.where(p >= dtype(1.0), infinity, x))
-  return x_nan_replaced
+  x_fix_boundaries = jnp.where(
+      p == dtype(0.0), -infinity, jnp.where(p == dtype(1.0), infinity, x))
+  return x_fix_boundaries
 
 
 @partial(custom_derivatives.custom_jvp, nondiff_argnums=(1,))


Test patch:
diff --git a/tests/lax_scipy_special_functions_test.py b/tests/lax_scipy_special_functions_test.py
--- a/tests/lax_scipy_special_functions_test.py
+++ b/tests/lax_scipy_special_functions_test.py
@@ -114,7 +114,7 @@ def op_record(name, nargs, dtypes, rng_factory, test_grad, nondiff_argnums=(), t
     ),
     op_record(
         "ndtri", 1, float_dtypes,
-        functools.partial(jtu.rand_uniform, low=0.05, high=0.95), True,
+        functools.partial(jtu.rand_uniform, low=0.0, high=1.0), True,
     ),
     op_record(
         "ndtr", 1, float_dtypes, jtu.rand_default, True
@@ -218,6 +218,13 @@ def testGammaSign(self):
     self._CheckAgainstNumpy(osp_special.gamma, lsp_special.gamma, args_maker, rtol=rtol)
     self._CompileAndCheck(lsp_special.gamma, args_maker, rtol=rtol)
 
+  def testNdtriExtremeValues(self):
+    # Testing at the extreme values (bounds (0. and 1.) and outside the bounds).
+    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.ndtri, lsp_special.ndtri, args_maker, rtol=rtol)
+    self._CompileAndCheck(lsp_special.ndtri, args_maker, rtol=rtol)
 
 
 if __name__ == "__main__":