Repository: google/jax

Issue:
jax.numpy.nanpercentile with axis as tuple
currently, the following code results in a NotImplementedError:
```
jax.numpy.nanpercentile(jnp.ones((3,3,12,24)), 99.9, axis=(0,1,2))
```

```
 File "/home/clemens/.local/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 6359, in _quantile                                       
    raise NotImplementedError("Tuple values for axis are not implemented")   
```
However the docstring actually explicitly mentions tuple as a valid option for axis: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.nanpercentile.html

same for jax.numpy.percentile.html


Discussion:
Thanks for the report – this feature hasn't been implemented yet. Regarding the docstring, the bulk of the content is copied directly from the corresponding numpy function (it's easy to miss, but at the top you'll see something like "original docstring below") because of that, the details are not always accurate, particularly when corner cases are explicitly unimplemented.

We can leave this issue open to track implementing this feature.
I'd like to give this one a try. Might take some time though as I'm quite new to Jax. If that's no problem, feel free to assign it to me!

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
@@ -6406,13 +6406,31 @@ def _quantile(a, q, axis, interpolation, keepdims, squash_nans):
     raise ValueError("interpolation can only be 'linear', 'lower', 'higher', "
                      "'midpoint', or 'nearest'")
   a, q = _promote_dtypes_inexact(a, q)
+  keepdim = []
   if issubdtype(a.dtype, np.complexfloating):
     raise ValueError("quantile does not support complex input, as the operation is poorly defined.")
   if axis is None:
     a = ravel(a)
     axis = 0
   elif isinstance(axis, tuple):
-    raise NotImplementedError("Tuple values for axis are not implemented")
+    keepdim = list(shape(a))
+    nd = ndim(a)
+    axis = tuple([_canonicalize_axis(ax, nd) for ax in axis])
+    if len(set(axis)) != len(axis):
+        raise ValueError('repeated axis')
+    for ax in axis:
+      keepdim[ax] = 1
+
+    keep = set(range(nd)) - set(axis)
+    # prepare permutation
+    dimensions = list(range(nd))
+    for i, s in enumerate(sorted(keep)):
+      dimensions[i], dimensions[s] = dimensions[s], dimensions[i]
+    do_not_touch_shape = tuple(x for idx,x in enumerate(shape(a)) if idx not in axis)
+    touch_shape = tuple(x for idx,x in enumerate(shape(a)) if idx in axis)
+    a = lax.reshape(a, do_not_touch_shape + (int(np.prod(touch_shape)),), dimensions)
+    keepdim = tuple(keepdim)
+    axis = _canonicalize_axis(-1, ndim(a))
   else:
     axis = _canonicalize_axis(axis, ndim(a))
 
@@ -6499,7 +6517,10 @@ def _quantile(a, q, axis, interpolation, keepdims, squash_nans):
     result = lax.mul(lax.add(low_value, high_value), _constant_like(low_value, 0.5))
   else:
     raise ValueError(f"interpolation={interpolation!r} not recognized")
-
+  if keepdims and keepdim:
+    if q_ndim > 0:
+        keepdim = (shape(q)[0],) + keepdim
+    result = reshape(result,  keepdim)
   return lax.convert_element_type(result, a.dtype)
 
 


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
@@ -4595,10 +4595,14 @@ def args_maker(): return []
         for a_shape, axis in (
           ((7,), None),
           ((47, 7), 0),
+          ((47, 7), ()),
           ((4, 101), 1),
+          ((4, 47, 7), (1, 2)),
+          ((4, 47, 7), (0, 2)),
+          ((4, 47, 7), (1, 0, 2)),
         )
         for q_dtype in [np.float32]
-        for q_shape in scalar_shapes + [(4,)]
+        for q_shape in scalar_shapes + [(1,), (4,)]
         for keepdims in [False, True]
         for method in ['linear', 'lower', 'higher', 'nearest', 'midpoint']))
   def testQuantile(self, op, a_rng, q_rng, a_shape, a_dtype, q_shape, q_dtype,