【Hackathon 5th No.116】paddle.quantile/nanquantile 功能增强

python/paddle/tensor/stat.py

@@ -468,8 +468,8 @@ def _compute_quantile(x, q, axis=None, keepdim=False, ignore_nan=False):
 
     Args:
         x (Tensor): The input Tensor, it's data type can be float32, float64, int32, int64.
-        q (int|float|list): The q for calculate quantile, which should be in range [0, 1]. If q is a list,
-            each q will be calculated and the first dimension of output is same to the number of ``q`` .
+        q (int|float|list|Tensor): The q for calculate quantile, which should be in range [0, 1]. If q is a list,
+            a 1-D Tensor or a 0-D Tensor, each q will be calculated and the first dimension of output is same to the number of ``q`` .
         axis (int|list, optional): The axis along which to calculate quantile. ``axis`` should be int or list of int.
             ``axis`` should be in range [-D, D), where D is the dimensions of ``x`` .
             If ``axis`` is less than 0, it works the same way as :math:`axis + D`.
@@ -498,8 +498,15 @@ def _compute_quantile(x, q, axis=None, keepdim=False, ignore_nan=False):
     elif isinstance(q, (list, tuple)):
         if len(q) <= 0:
             raise ValueError("q should not be empty")
+    elif isinstance(q, paddle.Tensor):
+        if len(q.shape) > 1:
+            raise ValueError("q should be a 0-D tensor or a 1-D tensor")
+        elif len(q.shape) == 0:
+            q = [q]
     else:
-        raise TypeError("Type of q should be int, float, list or tuple.")
+        raise TypeError(
+            "Type of q should be int, float, list or tuple, or tensor"
+        )
 
     # Validate axis
     dims = len(x.shape)
@@ -603,8 +610,8 @@ def quantile(x, q, axis=None, keepdim=False):
 
     Args:
         x (Tensor): The input Tensor, it's data type can be float32, float64, int32, int64.
-        q (int|float|list): The q for calculate quantile, which should be in range [0, 1]. If q is a list,
-            each q will be calculated and the first dimension of output is same to the number of ``q`` .
+        q (int|float|list|Tensor): The q for calculate quantile, which should be in range [0, 1]. If q is a list,
+            a 1-D Tensor or a 0-D Tensor, each q will be calculated and the first dimension of output is same to the number of ``q`` .
         axis (int|list, optional): The axis along which to calculate quantile. ``axis`` should be int or list of int.
             ``axis`` should be in range [-D, D), where D is the dimensions of ``x`` .
             If ``axis`` is less than 0, it works the same way as :math:`axis + D`.
@@ -671,8 +678,8 @@ def nanquantile(x, q, axis=None, keepdim=False):
 
     Args:
         x (Tensor): The input Tensor, it's data type can be float32, float64, int32, int64.
-        q (int|float|list): The q for calculate quantile, which should be in range [0, 1]. If q is a list,
-            each q will be calculated and the first dimension of output is same to the number of ``q`` .
+        q (int|float|list|Tensor): The q for calculate quantile, which should be in range [0, 1]. If q is a list or
+            a 1-D Tensor, each q will be calculated and the first dimension of output is same to the number of ``q`` .
         axis (int|list, optional): The axis along which to calculate quantile. ``axis`` should be int or list of int.
             ``axis`` should be in range [-D, D), where D is the dimensions of ``x`` .
             If ``axis`` is less than 0, it works the same way as :math:`axis + D`.

test/legacy_test/test_quantile_and_nanquantile.py

@@ -118,6 +118,41 @@ def test_nanquantile_all_NaN(self):
             paddle_res.numpy(), np_res, rtol=1e-05, equal_nan=True
         )
 
+    def test_backward(self):
+        def check_grad(x, q, axis, target_gard, apis=None):
+            for op, _ in apis or API_list:
+                x = paddle.to_tensor(x, dtype='float32', stop_gradient=False)
+                op(x, q, axis).sum().backward()
+                np.testing.assert_allclose(
+                    x.grad.numpy(),
+                    np.array(target_gard, dtype='float32'),
+                    rtol=1e-05,
+                    equal_nan=True,
+                )
+
+        check_grad([1, 2, 3], 0.5, 0, [0, 1, 0])
+        check_grad(
+            [1, 2, 3, 4] * 2, [0.55, 0.7], 0, [0, 0, 0.95, 0, 0, 0.15, 0.9, 0]
+        )
+        check_grad(
+            [[1, 2, 3], [4, 5, 6]],
+            [0.3, 0.7],
+            1,
+            [[0.4, 1.2, 0.4], [0.4, 1.2, 0.4]],
+        )
+        # quantile
+        check_grad(
+            [1, float('nan'), 3], 0.5, 0, [0, 1, 0], [(paddle.quantile, None)]
+        )
+        # nanquantile
+        check_grad(
+            [1, float('nan'), 3],
+            0.5,
+            0,
+            [0.5, 0, 0.5],
+            [(paddle.nanquantile, None)],
+        )
+
 
 class TestMuitlpleQ(unittest.TestCase):
     """
@@ -149,6 +184,24 @@ def test_quantile_multiple_axis_keepdim(self):
         )
         np.testing.assert_allclose(paddle_res.numpy(), np_res, rtol=1e-05)
 
+    def test_quantile_with_tensor_input(self):
+        x = paddle.to_tensor(self.input_data)
+        paddle_res = paddle.quantile(
+            x, q=paddle.to_tensor([0.1, 0.2]), axis=[1, 2], keepdim=True
+        )
+        np_res = np.quantile(
+            self.input_data, q=[0.1, 0.2], axis=[1, 2], keepdims=True
+        )
+        np.testing.assert_allclose(paddle_res.numpy(), np_res, rtol=1e-05)
+
+    def test_quantile_with_zero_dim_tensor_input(self):
+        x = paddle.to_tensor(self.input_data)
+        paddle_res = paddle.quantile(
+            x, q=paddle.to_tensor(0.1), axis=[1, 2], keepdim=True
+        )
+        np_res = np.quantile(self.input_data, q=0.1, axis=[1, 2], keepdims=True)
+        np.testing.assert_allclose(paddle_res.numpy(), np_res, rtol=1e-05)
+
 
 class TestError(unittest.TestCase):
     """
@@ -209,6 +262,19 @@ def test_axis_value_error_2():
 
         self.assertRaises(ValueError, test_axis_value_error_2)
 
+        # Test error when q is not a 1-D tensor
+        def test_tensor_input_1():
+            paddle_res = paddle.quantile(
+                self.x, q=paddle.randn((2, 3)), axis=[1, -10]
+            )
+
+        self.assertRaises(ValueError, test_tensor_input_1)
+
+        def test_type_q():
+            paddle_res = paddle.quantile(self.x, q={1}, axis=[1, -10])
+
+        self.assertRaises(TypeError, test_type_q)
+
 
 class TestQuantileRuntime(unittest.TestCase):
     """