Repository: google/jax

Issue:
Add all interpolation methods to quantile (and percentile)
# Description

At the moment, only the "linear" interpolation method from [`numpy.quantile`](https://docs.scipy.org/doc/numpy/reference/generated/numpy.quantile.html) and by extension [`numpy.percentile`](https://docs.scipy.org/doc/numpy/reference/generated/numpy.percentile.html) are implemented (as [noted in Issue #70](https://github.com/google/jax/issues/70#issuecomment-509322341)) 

https://github.com/google/jax/blob/1cf708ea77fae0fc1418e0b944f0115646f7f2ec/jax/numpy/lax_numpy.py#L3446-L3447

It would be useful to also have the "nearest" interpolation method, along with the others ("lower", "higher", "midpoint") implemented as well.

cc @lukasheinrich @kratsg


Discussion:
Agreed, this would be nice to have. Contributions would be very welcome here!
> Agreed, this would be nice to have

Great to hear that you're in agreement. I realistically won't have time to give this a shot for at least 2 weeks, but I'll try to do so then and open a PR once I pass the tests. If someone beats me to it that's great too of course. :)

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -3443,12 +3443,12 @@ def quantile(a, q, axis=None, out=None, overwrite_input=False,
     msg = ("jax.numpy.quantile does not support overwrite_input=True or "
            "out != None")
     raise ValueError(msg)
-  if interpolation != "linear":
-    raise NotImplementedError("Only interpolation='linear' is implemented")
-  return _quantile(a, q, axis, keepdims)
+  if interpolation not in ["linear", "lower", "higher", "midpoint", "nearest"]:
+    raise ValueError("interpolation can only be 'linear', 'lower', 'higher', 'midpoint', or 'nearest'")
+  return _quantile(a, q, axis, interpolation, keepdims)
 
-@partial(jit, static_argnums=(2, 3))
-def _quantile(a, q, axis, keepdims):
+@partial(jit, static_argnums=(2, 3, 4))
+def _quantile(a, q, axis, interpolation, keepdims):
   a = asarray(a)
   if axis is None:
     a = ravel(a)
@@ -3477,7 +3477,7 @@ def _quantile(a, q, axis, keepdims):
   n = a_shape[axis]
   q = lax.mul(q, _constant_like(q, n - 1))
   low = lax.floor(q)
-  high = lax.add(low, _constant_like(low, 1))
+  high = lax.ceil(q)
   high_weight = lax.sub(q, low)
   low_weight = lax.sub(_constant_like(high_weight, 1), high_weight)
 
@@ -3506,9 +3506,23 @@ def _quantile(a, q, axis, keepdims):
                                       broadcast_dimensions=(0,))
     high_weight = lax.broadcast_in_dim(high_weight, high_value.shape,
                                       broadcast_dimensions=(0,))
-  return lax.convert_element_type(
-    lax.add(lax.mul(low_value.astype(q.dtype), low_weight),
-            lax.mul(high_value.astype(q.dtype), high_weight)), a.dtype)
+
+  if interpolation == "linear":
+    result = lax.add(lax.mul(low_value.astype(q.dtype), low_weight),
+                     lax.mul(high_value.astype(q.dtype), high_weight))
+  elif interpolation == "lower":
+    result = low_value
+  elif interpolation == "higher":
+    result = high_value
+  elif interpolation == "nearest":
+    pred = lax.le(high_weight, _constant_like(high_weight, 0.5))
+    result = lax.select(pred, low_value, high_value)
+  elif interpolation == "midpoint":
+    result = lax.mul(lax.add(low_value, high_value), _constant_like(low_value, 0.5))
+  else:
+    raise ValueError(f"interpolation={interpolation!r} not recognized")
+
+  return lax.convert_element_type(result, a.dtype)
 
 
 @_wraps(onp.percentile)


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
@@ -2093,19 +2093,19 @@ def testIx_(self, rng_factory, shapes, dtypes):
 
   @parameterized.named_parameters(jtu.cases_from_list(
         {"testcase_name":
-           "_op={}_a_shape={}_q_shape={}_axis={}_keepdims={}".format(
+           "_op={}_a_shape={}_q_shape={}_axis={}_keepdims={}_interpolation={}".format(
              op,
              jtu.format_shape_dtype_string(a_shape, a_dtype),
              jtu.format_shape_dtype_string(q_shape, q_dtype),
-             axis, keepdims),
+             axis, keepdims, interpolation),
          "a_rng": jtu.rand_default(), "q_rng": q_rng, "op": op,
          "a_shape": a_shape, "a_dtype": a_dtype,
          "q_shape": q_shape, "q_dtype": q_dtype, "axis": axis,
-         "keepdims": keepdims}
+         "keepdims": keepdims,
+         "interpolation": interpolation}
         for (op, q_rng) in (
           ("percentile", jtu.rand_uniform(low=0., high=100.)),
           ("quantile", jtu.rand_uniform(low=0., high=1.)),
-          ("median", jtu.rand_uniform(low=0., high=1.)),
         )
         for a_dtype in float_dtypes
         for a_shape, axis in (
@@ -2115,20 +2115,20 @@ def testIx_(self, rng_factory, shapes, dtypes):
         )
         for q_dtype in [onp.float32]
         for q_shape in scalar_shapes + [(4,)]
-        for keepdims in [False, True]))
+        for keepdims in [False, True]
+        for interpolation in ['linear', 'lower', 'higher', 'nearest', 'midpoint']))
   def testQuantile(self, op, a_rng, q_rng, a_shape, a_dtype, q_shape, q_dtype,
-                   axis, keepdims):
+                   axis, keepdims, interpolation):
     if op == "quantile" and numpy_version < (1, 15):
       raise SkipTest("Numpy < 1.15 does not have np.quantile")
-    if op == "median":
-        args_maker = lambda: [a_rng(a_shape, a_dtype)]
-    else:
-        args_maker = lambda: [a_rng(a_shape, a_dtype), q_rng(q_shape, q_dtype)]
+    args_maker = lambda: [a_rng(a_shape, a_dtype), q_rng(q_shape, q_dtype)]
     def onp_fun(*args):
       args = [x if jnp.result_type(x) != jnp.bfloat16 else
               onp.asarray(x, onp.float32) for x in args]
-      return getattr(onp, op)(*args, axis=axis, keepdims=keepdims)
-    jnp_fun = partial(getattr(jnp, op), axis=axis, keepdims=keepdims)
+      return getattr(onp, op)(*args, axis=axis, keepdims=keepdims,
+                     interpolation=interpolation)
+    jnp_fun = partial(getattr(jnp, op), axis=axis, keepdims=keepdims,
+                      interpolation=interpolation)
     # TODO(phawkins): we currently set dtype=False because we aren't as
     # aggressive about promoting to float64. It's not clear we want to mimic
     # Numpy here.
@@ -2140,6 +2140,39 @@ def onp_fun(*args):
     self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True, rtol=tol)
 
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+        {"testcase_name":
+           "_op=median_a_shape={}_axis={}_keepdims={}".format(
+             jtu.format_shape_dtype_string(a_shape, a_dtype),
+             axis, keepdims),
+         "a_rng": jtu.rand_default(), 
+         "a_shape": a_shape, "a_dtype": a_dtype,
+         "axis": axis,
+         "keepdims": keepdims}
+        for a_dtype in float_dtypes
+        for a_shape, axis in (
+          ((7,), None),
+          ((47, 7), 0),
+          ((4, 101), 1),
+        )
+        for keepdims in [False, True]))
+  def testMedian(self, a_rng, a_shape, a_dtype, axis, keepdims):
+    args_maker = lambda: [a_rng(a_shape, a_dtype)]
+    def onp_fun(*args):
+      args = [x if jnp.result_type(x) != jnp.bfloat16 else
+              onp.asarray(x, onp.float32) for x in args]
+      return onp.median(*args, axis=axis, keepdims=keepdims)
+    jnp_fun = partial(jnp.median, axis=axis, keepdims=keepdims)
+    # TODO(phawkins): we currently set dtype=False because we aren't as
+    # aggressive about promoting to float64. It's not clear we want to mimic
+    # Numpy here.
+    tol_spec = {onp.float32: 2e-4, onp.float64: 5e-6}
+    tol = jtu.tolerance(a_dtype, tol_spec)
+    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=False,
+                            tol=tol)
+    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True, rtol=tol)
+
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}".format(
           jtu.format_shape_dtype_string(shape, dtype)),