【PIR API adaptor No.253、310】Migrate cumulative_trapezoid，trapezoid into pir (#59481)

enkilee · web-flow · commit f6de2c7b152e · 2023-12-04T14:51:26.000+08:00
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -6437,9 +6437,16 @@ def _trapezoid(y, x=None, dx=None, axis=-1, mode='sum'):
 
     if not (x is None or dx is None):
         raise ValueError("Not permitted to specify both x and dx input args.")
-    if y.dtype not in [paddle.float16, paddle.float32, paddle.float64]:
+    if y.dtype not in [
+        paddle.float16,
+        paddle.float32,
+        paddle.float64,
+        paddle.base.core.DataType.FLOAT16,
+        paddle.base.core.DataType.FLOAT32,
+        paddle.base.core.DataType.FLOAT64,
+    ]:
         raise TypeError(
-            f"The data type of input must be Tensor, and dtype should be one of ['paddle.float16', 'paddle.float32', 'paddle.float64'], but got {y.dtype}"
+            f"The data type of input must be Tensor, and dtype should be one of ['paddle.float16', 'paddle.float32', 'paddle.float64', 'paddle.base.core.DataType.FLOAT16', 'paddle.base.core.DataType.FLOAT32', 'paddle.base.core.DataType.FLOAT64'], but got {y.dtype}"
         )
 
     y_shape = y.shape
@@ -6453,9 +6460,16 @@ def _trapezoid(y, x=None, dx=None, axis=-1, mode='sum'):
         if dx.dim() > 1:
             raise ValueError(f'Expected dx to be a scalar, got dx={dx}')
     else:
-        if x.dtype not in [paddle.float16, paddle.float32, paddle.float64]:
+        if x.dtype not in [
+            paddle.float16,
+            paddle.float32,
+            paddle.float64,
+            paddle.base.core.DataType.FLOAT16,
+            paddle.base.core.DataType.FLOAT32,
+            paddle.base.core.DataType.FLOAT64,
+        ]:
             raise TypeError(
-                f"The data type of input must be Tensor, and dtype should be one of ['paddle.float16', 'paddle.float32', 'paddle.float64'], but got {x.dtype}"
+                f"The data type of input must be Tensor, and dtype should be one of ['paddle.float16', 'paddle.float32', 'paddle.float64', 'paddle.base.core.DataType.FLOAT16', 'paddle.base.core.DataType.FLOAT32', 'paddle.base.core.DataType.FLOAT64'], but got {x.dtype}"
             )
         # Reshape to correct shape
         if x.dim() == 1:
diff --git a/test/legacy_test/test_trapezoid.py b/test/legacy_test/test_trapezoid.py
@@ -17,6 +17,7 @@
 import numpy as np
 
 import paddle
+from paddle.pir_utils import test_with_pir_api
 
 
 class TestTrapezoidAPI(unittest.TestCase):
@@ -63,6 +64,7 @@ def test_dygraph(self):
         self.setUp()
         self.func_dygraph()
 
+    @test_with_pir_api
     def test_static(self):
         paddle.enable_static()
         places = [paddle.CPUPlace()]
@@ -162,6 +164,7 @@ class TestTrapezoidError(unittest.TestCase):
     def set_api(self):
         self.paddle_api = paddle.trapezoid
 
+    @test_with_pir_api
     def test_errors(self):
         self.set_api()
         with paddle.static.program_guard(
@@ -226,6 +229,7 @@ def set_api(self):
         self.paddle_api = paddle.trapezoid
         self.ref_api = np.trapz
 
+    @test_with_pir_api
     def test_fp16_with_gpu(self):
         paddle.enable_static()
         if paddle.base.core.is_compiled_with_cuda():
