Repository: google/jax

Issue:
Floating point deviation in `jax.numpy.percentile` with linear interpolation between `v0.2.20` and `v0.2.21`
Hi. There is some (very minor) deviations in the output of `jax.numpy.percentile` between `jax` [`v0.2.20` and `v0.2.21`](https://github.com/google/jax/compare/jax-v0.2.20...jax-v0.2.21) in the case that linear interpolation is used (the default). Interestingly, it is really in `jax.numpy.percentile` and not in `jax.numpy.quantile` as can be shown in the included example (for convenience this Issue also exists as [a GitHub Gist](https://gist.github.com/matthewfeickert/d8364201ddad6653315a62dc9b921318)).

## Minimal failing example

```python
# example.py
import jax
import jax.numpy as jnp
import numpy as np
from jax.config import config

config.update("jax_enable_x64", True)

if __name__ == "__main__":

    # percentile interpolation options:
    # This optional parameter specifies the interpolation method to use when the desired percentile lies between two data points i < j:
    # * ’linear’: i + (j - i) * fraction, where fraction is the fractional part of the index surrounded by i and j.
    # * ’lower’: i.
    # * ’higher’: j.
    # * ’nearest’: i or j, whichever is nearest.
    # * ’midpoint’: (i + j) / 2.

    input = [[10, 7, 4], [3, 2, 1]]
    print(f"input list: {input}")
    print(f"input list ravel: {np.asarray(input).ravel()}")
    # [10  7  4  3  2  1]

    print(f"\nNumPy v{np.__version__}")
    print(f"JAX v{jax.__version__}\n")
    numpy_array = np.asarray(input)
    print(f"{numpy_array=}")
    jax_array = jnp.asarray(input, dtype="float")
    print(f"{jax_array=}")

    print("\n# Checking quantile\n")
    assert np.quantile(numpy_array, 0) == 1.0
    assert np.quantile(numpy_array, 0.50) == 3.5
    assert np.quantile(numpy_array, 1) == 10
    assert np.quantile(numpy_array, 0.50, axis=1).tolist() == [7.0, 2.0]

    assert np.quantile(numpy_array, 0.50, interpolation="linear") == 3.5
    assert np.quantile(numpy_array, 0.50, interpolation="nearest") == 3.0
    assert np.quantile(numpy_array, 0.50, interpolation="lower") == 3.0
    assert np.quantile(numpy_array, 0.50, interpolation="midpoint") == 3.5
    assert np.quantile(numpy_array, 0.50, interpolation="higher") == 4.0

    assert jnp.quantile(jax_array, 0) == 1.0
    assert jnp.quantile(jax_array, 0.50) == 3.5
    assert jnp.quantile(jax_array, 1) == 10
    assert jnp.quantile(jax_array, 0.50, axis=1).tolist() == [7.0, 2.0]

    assert jnp.quantile(jax_array, 0.50, interpolation="linear") == 3.5
    assert jnp.quantile(jax_array, 0.50, interpolation="nearest") == 3.0
    assert jnp.quantile(jax_array, 0.50, interpolation="lower") == 3.0
    assert jnp.quantile(jax_array, 0.50, interpolation="midpoint") == 3.5
    assert jnp.quantile(jax_array, 0.50, interpolation="higher") == 4.0

    print("# Checking percentile")
    assert np.percentile(numpy_array, 0) == 1.0
    assert np.percentile(numpy_array, 50) == 3.5
    assert np.percentile(numpy_array, 100) == 10
    assert np.percentile(numpy_array, 50, axis=1).tolist() == [7.0, 2.0]

    assert np.percentile(numpy_array, 50, interpolation="linear") == 3.5
    assert np.percentile(numpy_array, 50, interpolation="nearest") == 3.0
    assert np.percentile(numpy_array, 50, interpolation="lower") == 3.0
    assert np.percentile(numpy_array, 50, interpolation="midpoint") == 3.5
    assert np.percentile(numpy_array, 50, interpolation="higher") == 4.0

    # default interpolation method is "linear"
    assert jnp.percentile(jax_array, 0) == 1.0
    assert jnp.percentile(jax_array, 50) == 3.5  # 3.499999761581421
    assert jnp.percentile(jax_array, 100) == 10  # 9.999998092651367
    assert jnp.percentile(jax_array, 50, axis=1).tolist() == [7.0, 2.0]

    assert jnp.percentile(jax_array, 50, interpolation="linear") == 3.5  # 3.499999761581421
    assert jnp.percentile(jax_array, 50, interpolation="nearest") == 3.0
    assert jnp.percentile(jax_array, 50, interpolation="lower") == 3.0
    assert jnp.percentile(jax_array, 50, interpolation="midpoint") == 3.5
    assert jnp.percentile(jax_array, 50, interpolation="higher") == 4.0
```

```console
user@machine:~$ python --version
Python 3.9.6
user@machine:~$ python -m venv /tmp/venv && . /tmp/venv/bin/activate
(venv) user@machine:~$ python -m pip install --upgrade pip setuptools wheel
(venv) user@machine:~$ cat requirements_passing.txt
jax==0.2.20
jaxlib==0.1.69
(venv) user@machine:~$ python -m pip install -r requirements_passing.txt
(venv) user@machine:~$ python example.py
input list: [[10, 7, 4], [3, 2, 1]]
input list ravel: [10  7  4  3  2  1]

NumPy v1.21.4
JAX v0.2.20

numpy_array=array([[10,  7,  4],
       [ 3,  2,  1]])
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
jax_array=DeviceArray([[10.,  7.,  4.],
             [ 3.,  2.,  1.]], dtype=float64)

# Checking quantile

# Checking percentile
(venv) user@machine:~$ cat requirements_failing.txt
jax==0.2.21
jaxlib==0.1.69
(venv) user@machine:~$ python -m pip install -r requirements_failing.txt
(venv) user@machine:~$ python example.py
input list: [[10, 7, 4], [3, 2, 1]]
input list ravel: [10  7  4  3  2  1]

NumPy v1.21.4
JAX v0.2.21

numpy_array=array([[10,  7,  4],
       [ 3,  2,  1]])
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
jax_array=DeviceArray([[10.,  7.,  4.],
             [ 3.,  2.,  1.]], dtype=float64)

# Checking quantile

# Checking percentile
Traceback (most recent call last):
  File "/home/feickert/Code/debug/jax-percentile-drift/example.py", line 67, in <module>
    assert jnp.percentile(jax_array, 50) == 3.5  # 3.499999761581421
AssertionError
```

## Notes

Comparing the code for [`v0.2.20`](https://github.com/google/jax/blob/jax-v0.2.20/jax/_src/numpy/lax_numpy.py#L5905-L5912)

https://github.com/google/jax/blob/a7b61c0e00d1b535df8a30a82edc0074884d5f4c/jax/_src/numpy/lax_numpy.py#L5905-L5912


and [`v0.2.21`](https://github.com/google/jax/blob/jax-v0.2.21/jax/_src/numpy/lax_numpy.py#L6420-L6429)

https://github.com/google/jax/blob/dbeb97d394740bfd122a46249c967139c10d3f11/jax/_src/numpy/lax_numpy.py#L6420-L6429

It seems (at first glance as I haven't dug into this yet) that the only relevant difference is the removal of `asarray(q)` in the `true_divide` call in PR #7747 (though I would think given the point of that PR that nothing should have changed)

```diff
-q = true_divide(asarray(q), float32(100.0))
+q = true_divide(q, float32(100.0))
```

This effect is quite minor, and probably poses no real significance in most cases, but it deviates from the docstring described behavior. Maybe the most obvious example is the extremes where the q-th percentile is 1 &mdash; which should return the element of the array object which is the maxima (in the example `10`) but instead returns the floating point approximation of that element (`9.999998092651367`).

## Request

Would it be possible to revert to the `v0.2.20` behavior? This would be more consistent with both the docstring and NumPy.

## JAX Issues checklist

Please:

- [x] Check for duplicate issues.
- [x] Provide a complete example of how to reproduce the bug, wrapped in triple backticks
- [x] If applicable, include full error messages/tracebacks.


Discussion:
I think the operative change here is that `jnp.percentile` is JIT-compiled by default starting in v0.2.21. Try running the following in version 0.2.20:
```python
import jax
import jax.numpy as jnp
import numpy as np
from jax.config import config

config.update("jax_enable_x64", True)

input = [[10, 7, 4], [3, 2, 1]]
numpy_array = np.asarray(input)
jax_array = jnp.asarray(input, dtype="float")

print("jax_version:", jax.__version__)
print("no jit:", jnp.percentile(jax_array, 50))
print("jit:", jax.jit(jnp.percentile)(jax_array, 50))
# jax_version: 0.2.20
# no jit: 3.5
# jit: 3.499999761581421
```
Why does JIT compiling cause this kind of inaccuracy? As part of compilation, XLA is free to re-arrange mathematical operations for efficiency, and sometimes this changes results slightly due to the imprecision inherent to floating point.

On further exploration, it looks like passing `np.float64(50)` rather than `50` fixes the issue. I think we could probably address this by being a bit more careful about dtype promotion within the `percentile` implementation.
> Why does JIT compiling cause this kind of inaccuracy? As part of compilation, XLA is free to re-arrange mathematical operations for efficiency, and sometimes this changes results slightly due to the imprecision inherent to floating point.

Thanks very much for the example and explanation @jakevdp &mdash; this is already quite helpful!

> I think we could probably address this by being a bit more careful about dtype promotion within the percentile implementation.

That would be great if possible in the future. :+1: 

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -6287,8 +6287,7 @@ def _quantile(a, q, axis, interpolation, keepdims, squash_nans):
   if interpolation not in ["linear", "lower", "higher", "midpoint", "nearest"]:
     raise ValueError("interpolation can only be 'linear', 'lower', 'higher', "
                      "'midpoint', or 'nearest'")
-  a = asarray(a, dtype=promote_types(_dtype(a), float32))
-  q = asarray(q, dtype=promote_types(_dtype(q), float32))
+  a, q = _promote_dtypes_inexact(a, q)
   if axis is None:
     a = ravel(a)
     axis = 0
@@ -6473,7 +6472,8 @@ def percentile(a, q, axis: Optional[Union[int, Tuple[int, ...]]] = None,
                out=None, overwrite_input=False, interpolation="linear",
                keepdims=False):
   _check_arraylike("percentile", a, q)
-  q = true_divide(q, float32(100.0))
+  a, q = _promote_dtypes_inexact(a, q)
+  q = true_divide(q, 100.0)
   return quantile(a, q, axis=axis, out=out, overwrite_input=overwrite_input,
                   interpolation=interpolation, keepdims=keepdims)
 


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -4416,6 +4416,12 @@ def np_fun(*args):
                             tol=tol)
     self._CompileAndCheck(jnp_fun, args_maker, rtol=tol)
 
+  @unittest.skipIf(not config.jax_enable_x64, "test requires X64")
+  @unittest.skipIf(jtu.device_under_test() != 'cpu', "test is for CPU float64 precision")
+  def testPercentilePrecision(self):
+    # Regression test for https://github.com/google/jax/issues/8513
+    x = jnp.float64([1, 2, 3, 4, 7, 10])
+    self.assertEqual(jnp.percentile(x, 50), 3.5)
 
   @parameterized.named_parameters(jtu.cases_from_list(
         {"testcase_name":