[PIR] Fix op mutable attribute use wrong dtype (#69375)

SigureMo · web-flow · commit 63d63dc04f85 · 2024-11-14T23:42:38.000+08:00
diff --git a/paddle/fluid/pir/dialect/op_generator/op_build_gen.py b/paddle/fluid/pir/dialect/op_generator/op_build_gen.py
@@ -542,7 +542,7 @@ def GenBuildOutputs(
                                   .attribute("value")
                                   .dyn_cast<paddle::dialect::ScalarAttribute>()
                                   .data()
-                                  .to<int>());
+                                  .to<{dtype}>());
   }}
   else {{
     {name} = phi::Scalar(-1);
diff --git a/paddle/fluid/pir/dialect/op_generator/op_infermeta_func_gen.py b/paddle/fluid/pir/dialect/op_generator/op_infermeta_func_gen.py
@@ -351,7 +351,7 @@ def GenBuildOutputsPart2(
                                   .attribute("value")
                                   .dyn_cast<paddle::dialect::ScalarAttribute>()
                                   .data()
-                                  .to<int>());
+                                  .to<{dtype}>());
   }}
   else {{
     {name} = phi::Scalar(-1);
diff --git a/paddle/phi/ops/yaml/inconsistent/update_ops.yaml b/paddle/phi/ops/yaml/inconsistent/update_ops.yaml
@@ -14,6 +14,5 @@
     param : [start, end, step]
     data_type : dtype
     backend : place
-  support_tensor : [start, end, step]
   interfaces : paddle::dialect::InferSymbolicShapeInterface
   traits : paddle::dialect::ForwardOnlyTrait
diff --git a/paddle/phi/ops/yaml/op_compat.yaml b/paddle/phi/ops/yaml/op_compat.yaml
@@ -239,6 +239,16 @@
     {start : Start, end : End, step : Step}
   outputs :
     out : Out
+  scalar:
+    start:
+      data_type : double
+      support_tensor : true
+    end:
+      data_type : double
+      support_tensor : true
+    step:
+      data_type : double
+      support_tensor : true
 
 - op : argmax(arg_max)
   inputs :
diff --git a/test/legacy_test/test_arange.py b/test/legacy_test/test_arange.py
@@ -229,5 +229,13 @@ def test_out(self):
         paddle.enable_static()
 
 
+class TestArangeStatic(unittest.TestCase):
+    def test_infermeta(self):
+        paddle.enable_static()
+        x = paddle.arange(0, 1 + 0.005, 0.005)
+        self.assertEqual(x.shape, [201])
+        paddle.disable_static()
+
+
 if __name__ == "__main__":
     unittest.main()
