draft version with some tests that are passing

Author: cakedev0

src/array_api_extra/__init__.py

@@ -11,6 +11,7 @@
     one_hot,
     pad,
     partition,
+    quantile,
     sinc,
 )
 from ._lib._at import at
@@ -48,6 +49,7 @@
     "one_hot",
     "pad",
     "partition",
+    "quantile",
     "setdiff1d",
     "sinc",
 ]

src/array_api_extra/_delegation.py

@@ -4,7 +4,7 @@
 from types import ModuleType
 from typing import Literal
 
-from ._lib import _funcs
+from ._lib import _funcs, _quantile
 from ._lib._utils._compat import (
     array_namespace,
     is_cupy_namespace,
@@ -930,18 +930,33 @@ def quantile(
     if ndim < 1:
         msg = "`a` must be at least 1-dimensional"
         raise TypeError(msg)
-    if (axis >= ndim) or (axis < -ndim):
+    if axis is not None and ((axis >= ndim) or (axis < -ndim)):
         message = "`axis` is not compatible with the dimension of `a`."
         raise ValueError(message)
 
+    # Array API states: Mixed integer and floating-point type promotion rules
+    # are not specified because behavior varies between implementations.
+    # => We choose to do:
+    dtype = (
+        xp.float64 if xp.isdtype(a.dtype, 'integral')
+        else xp.result_type(a, xp.asarray(q)) # both a and q are floats
+    )
+    device = get_device(a)
+    a = xp.asarray(a, dtype=dtype, device=device)
+    q = xp.asarray(q, dtype=dtype, device=device)
+
+    if xp.any((q > 1) | (q < 0) | xp.isnan(q)):
+        raise ValueError("`q` values must be in the range [0, 1]")
+
     # Delegate where possible.
-    if is_numpy_namespace(xp) or is_dask_namespace(xp):
+    if is_numpy_namespace(xp):
         return xp.quantile(a, q, axis=axis, method=method, keepdims=keepdims)
+    # No delegating for dask: I couldn't make it work
     is_linear = method == "linear"
     if (is_linear and is_jax_namespace(xp)) or is_cupy_namespace(xp):
         return xp.quantile(a, q, axis=axis, method=method, keepdims=keepdims)
     if is_linear and is_torch_namespace(xp):
         return xp.quantile(a, q, dim=axis, interpolation=method, keepdim=keepdims)
 
     # Otherwise call our implementation (will sort data)
-    return _funcs.quantile(a, q, axis=axis, method=method, keepdims=keepdims, xp=xp)
+    return _quantile.quantile(a, q, axis=axis, method=method, keepdims=keepdims, xp=xp)

src/array_api_extra/_lib/_quantile.py

@@ -17,7 +17,7 @@ def quantile(  # numpydoc ignore=PR01,RT01
 ):
     """See docstring in `array_api_extra._delegation.py`."""
     device = get_device(a)
-    floating_dtype = xp.result_type(a, xp.asarray(q))
+    floating_dtype = xp.float64 #xp.result_type(a, xp.asarray(q))
     a = xp.asarray(a, dtype=floating_dtype, device=device)
     q = xp.asarray(q, dtype=floating_dtype, device=device)
 
@@ -29,12 +29,13 @@ def quantile(  # numpydoc ignore=PR01,RT01
         q = xp.reshape(q, (1,))
 
     axis_none = axis is None
+    a_ndim = a.ndim
     if axis_none:
         a = xp.reshape(a, (-1,))
         axis = 0
     axis = int(axis)
 
-    n = eager_shape(a, axis)
+    n, = eager_shape(a, axis)
     # If data has length zero along `axis`, the result will be an array of NaNs just
     # as if the data had length 1 along axis and were filled with NaNs.
     if n == 0:
@@ -49,12 +50,12 @@ def quantile(  # numpydoc ignore=PR01,RT01
     # The hard part will be dealing with 0-weights and NaNs
     # But maybe a proper use of searchsorted + left/right side will work?
 
-    res = _quantile_hf(a, q, n, axis, xp)
+    res = _quantile_hf(a, q, float(n), axis, xp)
 
     # reshaping to conform to doc/other libs' behavior
     if axis_none:
         if keepdims:
-            res = xp.reshape(res, q.shape + (1,) * a.ndim)
+            res = xp.reshape(res, q.shape + (1,) * a_ndim)
     else:
         res = xp.moveaxis(res, axis, 0)
         if keepdims:
@@ -69,13 +70,18 @@ def quantile(  # numpydoc ignore=PR01,RT01
 def _quantile_hf(y: Array, p: Array, n: int, axis: int, xp: ModuleType):
     m = 1 - p
     jg = p*n + m - 1
+
     j = jg // 1
-    g = jg % 1
-    g[j < 0] = 0
     j = xp.clip(j, 0., n - 1)
     jp1 = xp.clip(j + 1, 0., n - 1)
     # `̀j` and `jp1` are 1d arrays
 
+    g = jg % 1
+    g = xp.where(j < 0, 0, g)  # equiv to g[j < 0] = 0, but work with strictest
+    new_g_shape = [1] * y.ndim
+    new_g_shape[axis] = g.shape[0]
+    g = xp.reshape(g, tuple(new_g_shape))
+
     return (
         (1 - g) * xp.take(y, xp.astype(j, xp.int64), axis=axis)
         + g * xp.take(y, xp.astype(jp1, xp.int64), axis=axis)

tests/test_funcs.py

@@ -29,6 +29,7 @@
     one_hot,
     pad,
     partition,
+    quantile,
     setdiff1d,
     sinc,
 )
@@ -1529,3 +1530,101 @@ def test_kind(self, xp: ModuleType, library: Backend):
         expected = xp.asarray([False, True, False, True])
         res = isin(a, b, kind="sort")
         xp_assert_equal(res, expected)
+
+
+@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no xp.take")
+class TestQuantile:
+    def test_basic(self, xp: ModuleType):
+        x = xp.asarray([1, 2, 3, 4, 5])
+        actual = quantile(x, 0.5)
+        expect = xp.asarray(3.0, dtype=xp.float64)
+        xp_assert_close(actual, expect)
+
+    def test_multiple_quantiles(self, xp: ModuleType):
+        x = xp.asarray([1, 2, 3, 4, 5])
+        actual = quantile(x, xp.asarray([0.25, 0.5, 0.75]))
+        expect = xp.asarray([2.0, 3.0, 4.0], dtype=xp.float64)
+        xp_assert_close(actual, expect)
+
+    def test_shape(self, xp: ModuleType):
+        a = xp.asarray(np.random.rand(3, 4, 5))
+        q = xp.asarray(np.random.rand(2))
+        assert quantile(a, q, axis=0).shape == (2, 4, 5)
+        assert quantile(a, q, axis=1).shape == (2, 3, 5)
+        assert quantile(a, q, axis=2).shape == (2, 3, 4)
+
+        assert quantile(a, q, axis=0, keepdims=True).shape == (2, 1, 4, 5)
+        assert quantile(a, q, axis=1, keepdims=True).shape == (2, 3, 1, 5)
+        assert quantile(a, q, axis=2, keepdims=True).shape == (2, 3, 4, 1)
+
+    def test_against_numpy(self, xp: ModuleType):
+        a_np = np.random.rand(3, 4, 5)
+        q_np = np.random.rand(2)
+        a = xp.asarray(a_np)
+        q = xp.asarray(q_np)
+        for keepdims in [False, True]:
+            for axis in [None, *range(a.ndim)]:
+                actual = quantile(a, q, axis=axis, keepdims=keepdims)
+                expected = np.quantile(a_np, q_np, axis=axis, keepdims=keepdims)
+                expected = xp.asarray(expected, dtype=xp.float64)
+                xp_assert_close(actual, expected, atol=1e-12)
+
+    def test_2d_axis(self, xp: ModuleType):
+        x = xp.asarray([[1, 2, 3], [4, 5, 6]])
+        actual = quantile(x, 0.5, axis=0)
+        expect = xp.asarray([2.5, 3.5, 4.5], dtype=xp.float64)
+        xp_assert_close(actual, expect)
+
+    def test_2d_axis_keepdims(self, xp: ModuleType):
+        x = xp.asarray([[1, 2, 3], [4, 5, 6]])
+        actual = quantile(x, 0.5, axis=0, keepdims=True)
+        expect = xp.asarray([[2.5, 3.5, 4.5]], dtype=xp.float64)
+        xp_assert_close(actual, expect)
+
+    def test_methods(self, xp: ModuleType):
+        x = xp.asarray([1, 2, 3, 4, 5])
+        methods = ["linear"] #"hazen", "weibull"]
+        for method in methods:
+            actual = quantile(x, 0.5, method=method)
+            # All methods should give reasonable results
+            assert 2.5 <= float(actual) <= 3.5
+
+    def test_edge_cases(self, xp: ModuleType):
+        x = xp.asarray([1, 2, 3, 4, 5])
+        # q = 0 should give minimum
+        actual = quantile(x, 0.0)
+        expect = xp.asarray(1.0, dtype=xp.float64)
+        xp_assert_close(actual, expect)
+
+        # q = 1 should give maximum
+        actual = quantile(x, 1.0)
+        expect = xp.asarray(5.0, dtype=xp.float64)
+        xp_assert_close(actual, expect)
+
+    def test_invalid_q(self, xp: ModuleType):
+        x = xp.asarray([1, 2, 3, 4, 5])
+        _ = quantile(x, 1.0)
+        # ^ FIXME: here just to make this test fail for sparse backend
+        # q > 1 should raise
+        with pytest.raises(
+            ValueError, match=r"`q` values must be in the range \[0, 1\]"
+        ):
+            _ = quantile(x, 1.5)
+        # q < 0 should raise
+        with pytest.raises(
+            ValueError, match=r"`q` values must be in the range \[0, 1\]"
+        ):
+            _ = quantile(x, -0.5)
+
+    def test_device(self, xp: ModuleType, device: Device):
+        if hasattr(device, 'type') and device.type == "meta":
+            pytest.xfail("No Tensor.item() on meta device")
+        x = xp.asarray([1, 2, 3, 4, 5], device=device)
+        actual = quantile(x, 0.5)
+        assert get_device(actual) == device
+
+    def test_xp(self, xp: ModuleType):
+        x = xp.asarray([1, 2, 3, 4, 5])
+        actual = quantile(x, 0.5, xp=xp)
+        expect = xp.asarray(3.0, dtype=xp.float64)
+        xp_assert_close(actual, expect)