ENH: implement quantile

Author: ev-br

autogen/numpy_api_dump.py

@@ -794,20 +794,6 @@ def putmask(a, mask, values):
     raise NotImplementedError
 
 
-def quantile(
-    a,
-    q,
-    axis=None,
-    out=None,
-    overwrite_input=False,
-    method="linear",
-    keepdims=False,
-    *,
-    interpolation=None,
-):
-    raise NotImplementedError
-
-
 def ravel(a, order="C"):
     raise NotImplementedError
 

torch_np/_detail/_reductions.py

@@ -228,10 +228,10 @@ def average(a_tensor, axis, w_tensor):
 
 
 
-def percentile(a_tensor, q_tensor, axis, method):
+def quantile(a_tensor, q_tensor, axis, method):
 
-    if (0 > q_tensor).any() or (q_tensor > 100).any():
-        raise ValueError("Percentiles must be in range [0, 100], got %s" % q_tensor)
+    if (0 > q_tensor).any() or (q_tensor > 1).any():
+        raise ValueError("Quantiles must be in range [0, 1], got %s" % q_tensor)
 
     if not a_tensor.dtype.is_floating_point:
         dtype = _scalar_types.default_float_type.torch_dtype
@@ -246,7 +246,7 @@ def percentile(a_tensor, q_tensor, axis, method):
         axis = _util.normalize_axis_tuple(axis, a_tensor.ndim)
     axis = _util.allow_only_single_axis(axis)
 
-    q_tensor = (q_tensor / 100.0).to(a_tensor.dtype)
+    q_tensor = q_tensor.to(a_tensor.dtype)
 
     (a_tensor, q_tensor), axis = _util.axis_none_ravel(a_tensor, q_tensor, axis=axis)
 

torch_np/_wrapper.py

@@ -712,16 +712,39 @@ def percentile(
         raise ValueError("'interpolation' argument is deprecated; use 'method' instead")
 
     a_tensor, q_tensor = _helpers.to_tensors(a, q)
-
-    result = _reductions.percentile(a_tensor, q_tensor, axis, method)
+    result = _reductions.quantile(a_tensor, q_tensor / 100., axis, method)
 
     # keepdims
     if keepdims:
         result = _util.apply_keepdims(result, axis, a_tensor.ndim)    
+    return _helpers.result_or_out(result, out, promote_scalar=True)
+
 
+def quantile(
+    a,
+    q,
+    axis=None,
+    out=None,
+    overwrite_input=False,
+    method="linear",
+    keepdims=False,
+    *,
+    interpolation=None,
+):
+    if interpolation is not None:
+        raise ValueError("'interpolation' argument is deprecated; use 'method' instead")
+
+    a_tensor, q_tensor = _helpers.to_tensors(a, q)
+    result = _reductions.quantile(a_tensor, q_tensor, axis, method)
+
+    # keepdims
+    if keepdims:
+        result = _util.apply_keepdims(result, axis, a_tensor.ndim)    
     return _helpers.result_or_out(result, out, promote_scalar=True)
 
 
+
+
 @asarray_replacer()
 def nanmean(a, axis=None, dtype=None, out=None, keepdims=NoValue, *, where=NoValue):
     if where is not NoValue:

torch_np/tests/numpy_tests/lib/test_function_base.py

@@ -3300,10 +3300,11 @@ def test_nan_q(self):
             np.percentile([1, 2, 3, 4.0], q)
 
 
-@pytest.mark.xfail(reason='TODO: implement')
+
 class TestQuantile:
     # most of this is already tested by TestPercentile
 
+    @pytest.mark.skip(reason="do not chase 1ulp")
     def test_max_ulp(self):
         x = [0.0, 0.2, 0.4]
         a = np.quantile(x, 0.45)
@@ -3318,6 +3319,7 @@ def test_basic(self):
         assert_equal(np.quantile(x, 1), 3.5)
         assert_equal(np.quantile(x, 0.5), 1.75)
 
+    @pytest.mark.xfail(reason="quantile w/integers or bools")
     def test_correct_quantile_value(self):
         a = np.array([True])
         tf_quant = np.quantile(True, False)
@@ -3328,6 +3330,7 @@ def test_correct_quantile_value(self):
         assert_array_equal(quant_res, a)
         assert_equal(quant_res.dtype, a.dtype)
 
+    @pytest.mark.skip(reason="support arrays of Fractions?")
     def test_fraction(self):
         # fractional input, integral quantile
         x = [Fraction(i, 2) for i in range(8)]
@@ -3355,10 +3358,9 @@ def test_fraction(self):
         x = np.arange(8)
         assert_equal(np.quantile(x, Fraction(1, 2)), Fraction(7, 2))
 
+    @pytest.mark.skip(reason="does not raise in numpy?")
     def test_complex(self):
         #See gh-22652
-        arr_c = np.array([0.5+3.0j, 2.1+0.5j, 1.6+2.3j], dtype='G')
-        assert_raises(TypeError, np.quantile, arr_c, 0.5)
         arr_c = np.array([0.5+3.0j, 2.1+0.5j, 1.6+2.3j], dtype='D')
         assert_raises(TypeError, np.quantile, arr_c, 0.5)
         arr_c = np.array([0.5+3.0j, 2.1+0.5j, 1.6+2.3j], dtype='F')
@@ -3376,12 +3378,14 @@ def test_no_p_overwrite(self):
         np.quantile(np.arange(100.), p, method="midpoint")
         assert_array_equal(p, p0)
 
+    @pytest.mark.xfail(reason="TODO: make quantile preserve integers")
     @pytest.mark.parametrize("dtype", np.typecodes["AllInteger"])
     def test_quantile_preserve_int_type(self, dtype):
         res = np.quantile(np.array([1, 2], dtype=dtype), [0.5],
                           method="nearest")
         assert res.dtype == dtype
 
+    @pytest.mark.xfail(reason="1) np.sort not implemented; 2) methods")
     @pytest.mark.parametrize("method",
              ['inverted_cdf', 'averaged_inverted_cdf', 'closest_observation',
               'interpolated_inverted_cdf', 'hazen', 'weibull', 'linear',