Repository: ivy-llc/ivy

Issue:
<nanpercentile>
Create nanpercentile frontend for jax numpy frontend


Discussion:

Resolving patch:
diff --git a/ivy/functional/frontends/jax/numpy/statistical.py b/ivy/functional/frontends/jax/numpy/statistical.py
--- a/ivy/functional/frontends/jax/numpy/statistical.py
+++ b/ivy/functional/frontends/jax/numpy/statistical.py
@@ -7,6 +7,7 @@
     handle_jax_dtype,
 )
 from ivy.functional.frontends.jax.numpy import promote_types_of_jax_inputs
+from ivy.functional.backends.jax.experimental.elementwise import _normalize_axis_tuple
 
 
 @to_ivy_arrays_and_back
@@ -359,6 +360,168 @@ def nanmin(
     return res.astype(ivy.dtype(a))
 
 
+@to_ivy_arrays_and_back
+@with_unsupported_dtypes(
+    {"0.4.14 and below": ("complex64", "complex128", "bfloat16", "bool", "float16")},
+    "jax",
+)
+def nanpercentile(
+    a, q, axis=None, out=None, overwrite_input=False, method="linear", keepdims=None
+):
+    def _remove_nan_1d(arr1d, overwrite_input=False):
+        if arr1d.dtype == object:
+            c = ivy.not_equal(arr1d, arr1d)
+        else:
+            c = ivy.isnan(arr1d)
+        s = ivy.nonzero(c)[0]
+        if s.size == arr1d.size:
+            return arr1d[:0], True
+        elif s.size == 0:
+            return arr1d, overwrite_input
+        else:
+            if not overwrite_input:
+                arr1d = arr1d.copy()
+
+                enonan = arr1d[-s.size :][~c[-s.size :]]
+                arr1d[s[: enonan.size]] = enonan
+
+                return arr1d[: -s.size], True
+
+    def _nanquantile_1d(arr1d, q, overwrite_input=False, method="linear"):
+        arr1d, overwrite_input = _remove_nan_1d(arr1d, overwrite_input=overwrite_input)
+        if arr1d.size == 0:
+            return ivy.full(q.shape, ivy.nan)
+        return ivy.quantile(arr1d, q, interpolation=method)
+
+    def apply_along_axis(func1d, axis, arr, *args, **kwargs):
+        ndim = ivy.get_num_dims(arr)
+        if axis is None:
+            raise ValueError("Axis must be an integer.")
+        if not -ndim <= axis < ndim:
+            raise ValueError(
+                f"axis {axis} is out of bounds for array of dimension {ndim}"
+            )
+        if axis < 0:
+            axis = axis + ndim
+
+        func = lambda elem: func1d(elem, *args, **kwargs)
+        for i in range(1, ndim - axis):
+            func = ivy.vmap(func, in_axes=i, out_axes=-1)
+        for i in range(axis):
+            func = ivy.vmap(func, in_axes=0, out_axes=0)
+
+        return ivy.asarray(func(arr))
+
+    def _nanquantile_ureduce_func(
+        a, q, axis=None, out=None, overwrite_input=False, method="linear"
+    ):
+        if axis is None or a.ndim == 1:
+            part = a.ravel()
+            result = _nanquantile_1d(
+                part, q, overwrite_input=overwrite_input, method=method
+            )
+        else:
+            result = apply_along_axis(
+                _nanquantile_1d, axis, a, q, overwrite_input, method
+            )
+
+            if q.ndim != 0:
+                result = ivy.moveaxis(result, axis, 0)
+
+        if out is not None:
+            out[...] = result
+
+        return result
+
+    def _ureduce(a, func, keepdims=False, **kwargs):
+        axis = kwargs.get("axis", None)
+        out = kwargs.get("out", None)
+
+        if keepdims is None:
+            keepdims = False
+
+        nd = a.ndim
+        if axis is not None:
+            axis = _normalize_axis_tuple(axis, nd)
+
+            if keepdims:
+                if out is not None:
+                    index_out = tuple(
+                        0 if i in axis else slice(None) for i in range(nd)
+                    )
+                    kwargs["out"] = out[(Ellipsis,) + index_out]
+
+            if len(axis) == 1:
+                kwargs["axis"] = axis[0]
+            else:
+                keep = set(range(nd)) - set(axis)
+                nkeep = len(keep)
+                # swap axis that should not be reduced to front
+                for i, s in enumerate(sorted(keep)):
+                    a = a.swapaxes(i, s)
+                # merge reduced axis
+                a = a.reshape(a.shape[:nkeep] + (-1,))
+                kwargs["axis"] = -1
+        else:
+            if keepdims:
+                if out is not None:
+                    index_out = (0,) * nd
+                    kwargs["out"] = out[(Ellipsis,) + index_out]
+
+        r = func(a, **kwargs)
+
+        if out is not None:
+            return out
+
+        if keepdims:
+            if axis is None:
+                index_r = (ivy.newaxis,) * nd
+            else:
+                index_r = tuple(
+                    ivy.newaxis if i in axis else slice(None) for i in range(nd)
+                )
+            r = r[(Ellipsis,) + index_r]
+
+        return r
+
+    def _nanquantile_unchecked(
+        a,
+        q,
+        axis=None,
+        out=None,
+        overwrite_input=False,
+        method="linear",
+        keepdims=None,
+    ):
+        """Assumes that q is in [0, 1], and is an ndarray."""
+        if a.size == 0:
+            return ivy.nanmean(a, axis, out=out, keepdims=keepdims)
+        return _ureduce(
+            a,
+            func=_nanquantile_ureduce_func,
+            q=q,
+            keepdims=keepdims,
+            axis=axis,
+            out=out,
+            overwrite_input=overwrite_input,
+            method=method,
+        )
+
+    a = ivy.array(a)
+    q = ivy.divide(q, 100.0)
+    q = ivy.array(q)
+    if q.ndim == 1 and q.size < 10:
+        for i in range(q.size):
+            if not (0.0 <= q[i] <= 1.0):
+                ivy.logging.warning("percentile s must be in the range [0, 100]")
+                return []
+    else:
+        if not (ivy.all(0 <= q) and ivy.all(q <= 1)):
+            ivy.logging.warning("percentile s must be in the range [0, 100]")
+            return []
+    return _nanquantile_unchecked(a, q, axis, out, overwrite_input, method, keepdims)
+
+
 @handle_jax_dtype
 @to_ivy_arrays_and_back
 def nanstd(


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_jax/test_numpy/test_statistical.py b/ivy_tests/test_ivy/test_frontends/test_jax/test_numpy/test_statistical.py
--- a/ivy_tests/test_ivy/test_frontends/test_jax/test_numpy/test_statistical.py
+++ b/ivy_tests/test_ivy/test_frontends/test_jax/test_numpy/test_statistical.py
@@ -192,6 +192,55 @@ def _get_dtype_value1_value2_cov(
     return [dtype], value1, value2, rowVar, bias, ddof, fweights, aweights
 
 
+@st.composite
+def _percentile_helper(draw):
+    large_abs_safety_factor = 2
+    small_abs_safety_factor = 2
+    dtype, values, axis = draw(
+        helpers.dtype_values_axis(
+            available_dtypes=helpers.get_dtypes("float"),
+            large_abs_safety_factor=large_abs_safety_factor,
+            small_abs_safety_factor=small_abs_safety_factor,
+            safety_factor_scale="log",
+            min_num_dims=1,
+            max_num_dims=5,
+            min_dim_size=2,
+            valid_axis=True,
+            allow_neg_axes=False,
+            min_axes_size=1,
+            force_int_axis=True,
+        )
+    )
+    q = draw(
+        st.one_of(
+            helpers.array_values(
+                dtype=helpers.get_dtypes("float"),
+                shape=helpers.get_shape(min_dim_size=1, max_num_dims=1, min_num_dims=1),
+                min_value=0.0,
+                max_value=100.0,
+                exclude_max=False,
+                exclude_min=False,
+            ),
+            st.floats(min_value=0.0, max_value=100.0),
+        )
+    )
+
+    interpolation_names = [
+        "linear",
+        "lower",
+        "higher",
+        "midpoint",
+        "nearest",
+    ]
+    interpolation = draw(
+        helpers.list_of_size(
+            x=st.sampled_from(interpolation_names),
+            size=1,
+        )
+    )
+    return dtype, values, axis, interpolation, q
+
+
 # --- Main --- #
 # ------------ #
 
@@ -1272,3 +1321,28 @@ def test_jax_var(
         atol=1e-3,
         rtol=1e-3,
     )
+
+
+@handle_frontend_test(
+    fn_tree="jax.numpy.nanpercentile",
+    dtype_and_x=_percentile_helper(),
+    keep_dims=st.booleans(),
+    test_gradients=st.just(False),
+    test_with_out=st.just(False),
+)
+def test_nanpercentile(
+    *, dtype_and_x, keep_dims, test_flags, backend_fw, fn_name, on_device
+):
+    input_dtype, x, axis, interpolation, q = dtype_and_x
+    helpers.test_function(
+        input_dtypes=input_dtype,
+        test_flags=test_flags,
+        backend_to_test=backend_fw,
+        fn_name=fn_name,
+        on_device=on_device,
+        a=x[0],
+        q=q,
+        axis=axis,
+        interpolation=interpolation[0],
+        keepdims=keep_dims,
+    )