add square inplace op (#69451)

NKNaN · web-flow · commit 5cfdddda7aeb · 2024-11-19T12:09:18.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc
@@ -186,7 +186,9 @@ OP_SAME_OPERANDS_AND_RESULT(LeakyRelu_)
 OP_SAME_OPERANDS_AND_RESULT(ThresholdedRelu)
 OP_SAME_OPERANDS_AND_RESULT(ThresholdedRelu_)
 OP_SAME_OPERANDS_AND_RESULT(SquareSr)
+OP_SAME_OPERANDS_AND_RESULT(SquareSr_)
 OP_SAME_OPERANDS_AND_RESULT(Square)
+OP_SAME_OPERANDS_AND_RESULT(Square_)
 OP_SAME_OPERANDS_AND_RESULT(Polygamma)
 OP_SAME_OPERANDS_AND_RESULT(Polygamma_)
 OP_SAME_OPERANDS_AND_RESULT(EnableCheckModelNanInf)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h
@@ -179,7 +179,9 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(ThresholdedRelu)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(ThresholdedRelu_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Selu)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SquareSr)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(SquareSr_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Square)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Square_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Polygamma)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Polygamma_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(EnableCheckModelNanInf)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -4720,6 +4720,7 @@
   kernel :
     func : square {dense -> dense},
            square_sr {selected_rows -> selected_rows}
+  inplace : (x -> out)
   backward : square_grad
   interfaces : paddle::dialect::InferSymbolicShapeInterface
 
diff --git a/python/paddle/tensor/__init__.py b/python/paddle/tensor/__init__.py
@@ -401,6 +401,7 @@
     sqrt,
     sqrt_,
     square,
+    square_,
     stanh,
     subtract,
     subtract_,
@@ -552,6 +553,7 @@
     'sqrt',
     'sqrt_',
     'square',
+    'square_',
     'stanh',
     'sum',
     'reduce_as',
diff --git a/test/legacy_test/test_inplace.py b/test/legacy_test/test_inplace.py
@@ -567,6 +567,18 @@ def inplace_api_processing(self, var):
         return var.rsqrt_()
 
 
+class TestDygraphInplaceSquare(TestDygraphInplace):
+    def init_data(self):
+        self.input_var_numpy = np.random.uniform(0, 5, [10, 20, 1])
+        self.dtype = "float32"
+
+    def non_inplace_api_processing(self, var):
+        return var.square()
+
+    def inplace_api_processing(self, var):
+        return var.square_()
+
+
 class TestDygraphInplaceClip(TestDygraphInplace):
     def non_inplace_api_processing(self, var):
         return var.clip(0.6, 1.5)
