ENH: implement median

Author: ev-br

autogen/numpy_api_dump.py

@@ -602,10 +602,6 @@ def may_share_memory(a, b, /, max_work=None):
     raise NotImplementedError
 
 
-def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):
-    raise NotImplementedError
-
-
 def meshgrid(*xi, copy=True, sparse=False, indexing="xy"):
     raise NotImplementedError
 

torch_np/_wrapper.py

@@ -708,16 +708,8 @@ def percentile(
     *,
     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 / 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)
+    return quantile(a, asarray(q) / 100., axis, out, overwrite_input, method,
+                    keepdims=keepdims)
 
 
 def quantile(
@@ -743,6 +735,9 @@ def quantile(
     return _helpers.result_or_out(result, out, promote_scalar=True)
 
 
+def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):
+    return quantile(a, 0.5, axis=axis, overwrite_input=overwrite_input,
+                    out=out, keepdims=keepdims)
 
 
 @asarray_replacer()

torch_np/tests/numpy_tests/lib/test_function_base.py

@@ -3474,7 +3474,7 @@ def test_linear_interpolation_formula_0d_inputs(self):
         assert nfb._lerp(a, b, t) == 2.6
 
 
-@pytest.mark.xfail(reason='TODO: implement')
+#@pytest.mark.xfail(reason='TODO: implement')
 class TestMedian:
 
     def test_basic(self):
@@ -3494,7 +3494,11 @@ def test_basic(self):
         assert_equal(a[0], np.median(a))
         a = np.array([0.0444502, 0.141249, 0.0463301])
         assert_equal(a[-1], np.median(a))
+
+    @pytest.mark.xfail(reason="median: scalar output vs 0-dim")
+    def test_basic_2(self):
         # check array scalar result
+        a = np.array([0.0444502, 0.141249, 0.0463301])
         assert_equal(np.median(a).ndim, 0)
         a[1] = np.nan
         assert_equal(np.median(a).ndim, 0)
@@ -3590,7 +3594,7 @@ def test_nan_behavior(self):
 
         # no axis
         assert_equal(np.median(a), np.nan)
-        assert_equal(np.median(a).ndim, 0)
+  #      assert_equal(np.median(a).ndim, 0)
 
         # axis0
         b = np.median(np.arange(24, dtype=float).reshape(2, 3, 4), 0)
@@ -3604,12 +3608,29 @@ def test_nan_behavior(self):
         b[1, 2] = np.nan
         assert_equal(np.median(a, 1), b)
 
+
+    @pytest.mark.xfail(reason="median: does not support tuple axes")
+    def test_nan_behavior_2(self):
+        a = np.arange(24, dtype=float).reshape(2, 3, 4)
+        a[1, 2, 3] = np.nan
+        a[1, 1, 2] = np.nan
+
         # axis02
         b = np.median(np.arange(24, dtype=float).reshape(2, 3, 4), (0, 2))
         b[1] = np.nan
         b[2] = np.nan
         assert_equal(np.median(a, (0, 2)), b)
 
+    @pytest.mark.xfail(reason="median: scalar vs 0-dim")
+    def test_nan_behavior_3(self):
+        a = np.arange(24, dtype=float).reshape(2, 3, 4)
+        a[1, 2, 3] = np.nan
+        a[1, 1, 2] = np.nan
+
+        # no axis
+        assert_equal(np.median(a).ndim, 0)
+
+    @pytest.mark.xfail(reason="median: torch.quantile does not handle empty tensors")
     @pytest.mark.skipif(IS_WASM, reason="fp errors don't work correctly")
     def test_empty(self):
         # mean(empty array) emits two warnings: empty slice and divide by 0
@@ -3640,6 +3661,7 @@ def test_empty(self):
             assert_equal(np.median(a, axis=2), b)
             assert_(w[0].category is RuntimeWarning)
 
+    @pytest.mark.xfail(reason="median: tuple axes not implemented")
     def test_extended_axis(self):
         o = np.random.normal(size=(71, 23))
         x = np.dstack([o] * 10)
@@ -3682,6 +3704,10 @@ def test_keepdims(self):
         d = np.ones((3, 5, 7, 11))
         assert_equal(np.median(d, axis=None, keepdims=True).shape,
                      (1, 1, 1, 1))
+
+    @pytest.mark.xfail(reason="median: tuple axis")
+    def test_keepdims_2(self):
+        d = np.ones((3, 5, 7, 11))
         assert_equal(np.median(d, axis=(0, 1), keepdims=True).shape,
                      (1, 1, 7, 11))
         assert_equal(np.median(d, axis=(0, 3), keepdims=True).shape,
@@ -3693,6 +3719,7 @@ def test_keepdims(self):
         assert_equal(np.median(d, axis=(0, 1, 3), keepdims=True).shape,
                      (1, 1, 7, 1))
 
+    @pytest.mark.xfail(reason="median: tuple axis")
     @pytest.mark.parametrize(
         argnames='axis',
         argvalues=[