【Hackthon 6th No. 10】Upgrade isinf to support int8 int16 uint8 -part (#63222)

NKNaN · web-flow · commit 09a522a95b3a · 2024-04-07T11:27:07.000+08:00
* upgrade isinf to support int8 int16 uint8

* update dtype
diff --git a/paddle/phi/kernels/cpu/isfinite_kernel.cc b/paddle/phi/kernels/cpu/isfinite_kernel.cc
@@ -27,7 +27,10 @@ PD_REGISTER_KERNEL(isinf,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
                    int,
-                   int64_t) {
+                   int64_t,
+                   int16_t,
+                   int8_t,
+                   uint8_t) {
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
 }
 
diff --git a/paddle/phi/kernels/gpu/isfinite_kernel.cu b/paddle/phi/kernels/gpu/isfinite_kernel.cu
@@ -27,7 +27,10 @@ PD_REGISTER_KERNEL(isinf,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
                    int,
-                   int64_t) {
+                   int64_t,
+                   int16_t,
+                   int8_t,
+                   uint8_t) {
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
 }
 
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -4400,7 +4400,7 @@ def isinf(x, name=None):
     Return whether every element of input tensor is `+/-INF` or not.
 
     Args:
-        x (Tensor): The input tensor, it's data type should be float16, float32, float64, int32, int64.
+        x (Tensor): The input tensor, it's data type should be float16, float32, float64, uint8, int8, int16, int32, int64.
         name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
 
     Returns:
@@ -4428,8 +4428,11 @@ def isinf(x, name=None):
                 'float16',
                 'float32',
                 'float64',
+                'int8',
+                'int16',
                 'int32',
                 'int64',
+                'uint8',
                 'uint16',
             ],
             'isinf',
diff --git a/test/legacy_test/test_isfinite_v2_op.py b/test/legacy_test/test_isfinite_v2_op.py
@@ -111,9 +111,33 @@ def np_data_generator(
     },
 ]
 
+TEST_META_DATA_ADDITIONAL = [
+    {
+        'low': 0.1,
+        'high': 1,
+        'np_shape': [2, 3, 4, 5],
+        'type': 'int8',
+        'sv_list': [np.inf, np.nan],
+    },
+    {
+        'low': 0,
+        'high': 100,
+        'np_shape': [11, 17, 10],
+        'type': 'int16',
+        'sv_list': [np.inf, np.nan],
+    },
+    {
+        'low': 0,
+        'high': 999,
+        'np_shape': [132],
+        'type': 'uint8',
+        'sv_list': [np.inf, np.nan],
+    },
+]
+
 
-def test(test_case, op_str, use_gpu=False):
-    for meta_data in TEST_META_DATA:
+def test(test_case, op_str, use_gpu=False, data_set=TEST_META_DATA):
+    for meta_data in data_set:
         meta_data = dict(meta_data)
         meta_data['op_str'] = op_str
         x_np, result_np = np_data_generator(**meta_data)
@@ -144,6 +168,9 @@ def test_nan(self):
     def test_finite(self):
         test(self, 'isfinite')
 
+    def test_inf_additional(self):
+        test(self, 'isinf', data_set=TEST_META_DATA_ADDITIONAL)
+
 
 class TestCUDANormal(unittest.TestCase):
     def test_inf(self):
@@ -155,6 +182,9 @@ def test_nan(self):
     def test_finite(self):
         test(self, 'isfinite', True)
 
+    def test_inf_additional(self):
+        test(self, 'isinf', True, data_set=TEST_META_DATA_ADDITIONAL)
+
 
 class TestError(unittest.TestCase):
     @test_with_pir_api
