[PIR] A-12 Adapt concat test_errors (#61036)

enkilee · web-flow · commit 9b87d22e4fc2 · 2024-02-21T18:58:10.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/utils/utils.cc b/paddle/fluid/pir/dialect/operator/utils/utils.cc
@@ -402,15 +402,33 @@ void CheckValueDataType(const pir::Value& value,
   DoValueCheck(value, input_name, expected_dtype, op_name);
 }
 
+bool IsSameDataTypeForValues(const std::vector<pir::Value>& vector_value) {
+  if (vector_value.size() <= 1) {
+    return true;
+  }
+  auto dtype = GetValueDataType(vector_value[0]);
+  for (size_t i = 1; i < vector_value.size(); ++i) {
+    if (GetValueDataType(vector_value[i]) != dtype) {
+      return false;
+    }
+  }
+  return true;
+}
+
 void CheckVectorOfValueDataType(const std::vector<pir::Value>& vector_value,
                                 const std::string& input_name,
                                 const std::string& op_name) {
   VLOG(6) << "CheckVectorOfValueDataType for " << op_name
           << ", input: " << input_name;
-  std::set<std::string> expected_dtype = GetRegisterDataType(op_name);
-  for (auto& value : vector_value) {
-    DoValueCheck(value, input_name, expected_dtype, op_name);
+  if (vector_value.size() == 0) {
+    return;
+  }
+  if (!IsSameDataTypeForValues(vector_value)) {
+    PADDLE_THROW(phi::errors::InvalidType(
+        "All the Values in the input must have the same data type."));
   }
+  std::set<std::string> expected_dtype = GetRegisterDataType(op_name);
+  DoValueCheck(vector_value[0], input_name, expected_dtype, op_name);
 }
 
 void CheckDataType(const phi::DataType& dtype,
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -1247,8 +1247,8 @@ def concat(x, axis=0, name=None):
     doesn't have any axis.
 
     Args:
-        x (list|tuple): ``x`` is a Tensor list or Tensor tuple which is with data type bool, float16,
-            float32, float64, int32, int64, int8, uint8. All the Tensors in ``x`` must have same data type.
+        x (list|tuple): ``x`` is a Tensor list or Tensor tuple which is with data type bool, float16, bfloat16,
+            float32, float64, int8, int16, int32, int64, uint8, uint16, complex64, complex128. All the Tensors in ``x`` must have same data type.
         axis (int|Tensor, optional): Specify the axis to operate on the input Tensors.
             Tt should be integer or 0-D int Tensor with shape []. The effective range is [-R, R), where R is Rank(x). When ``axis < 0``,
             it works the same way as ``axis+R``. Default is 0.
@@ -1319,13 +1319,17 @@ def concat(x, axis=0, name=None):
                     [
                         'bool',
                         'float16',
+                        'bfloat16',
                         'float32',
                         'float64',
+                        'int16',
                         'int32',
                         'int64',
                         'int8',
                         'unit8',
                         'uint16',
+                        'complex64',
+                        'complex128',
                     ],
                     'concat',
                 )
diff --git a/test/legacy_test/test_concat_op.py b/test/legacy_test/test_concat_op.py
@@ -940,5 +940,29 @@ def if_enable_cinn(self):
         pass
 
 
+class TestConcatOpErrorWithPir(unittest.TestCase):
+    @test_with_pir_api
+    def test_errors_with_pir(self):
+        paddle.enable_static()
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
+            # The type of axis in concat_op should be int or Variable.
+            x6 = paddle.static.data(shape=[-1, 4], dtype='float32', name='x6')
+            x7 = paddle.static.data(shape=[-1, 4], dtype='float32', name='x7')
+            x8 = paddle.static.data(shape=[-1, 4], dtype='float64', name='x8')
+
+            def test_axis_type():
+                paddle.concat([x6, x7], 3.2)
+
+            self.assertRaises(TypeError, test_axis_type)
+
+            # The input dtype must be same.
+            def test_input_same_dtype():
+                paddle.concat([x7, x8])
+
+            self.assertRaises(TypeError, test_input_same_dtype)
+
+
 if __name__ == '__main__':
     unittest.main()
