chore: import upstream snapshot with attribution

test/dygraph_to_static/test_python_op.py

@@ -0,0 +1,159 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+from collections.abc import Callable
+
+import numpy as np
+from dygraph_to_static_utils import (
+    Dy2StTestBase,
+    static_guard,
+)
+
+import paddle
+from paddle import Tensor
+from paddle.jit import sot
+
+
+@paddle.static.register_op(
+    name="fn_with_breakgraph",
+    infer_meta=lambda x, y: paddle.static.MetaTensor(
+        dtype=x.dtype, shape=x.shape
+    ),
+    input_names=["x", "y"],
+    output_names=["out"],
+    inplace_map={},
+)
+def fn_with_breakgraph(x: Tensor, y: Tensor) -> Tensor:
+    x = x + 1
+    sot.psdb.breakgraph()
+    y = y + 1
+    return x + y
+
+
+@paddle.static.register_op(
+    name="fn_with_numpy_operation",
+    infer_meta=lambda x, y: paddle.static.MetaTensor(
+        dtype=paddle.int32, shape=x.shape[:-1]
+    ),
+    input_names=["x", "y"],
+    output_names=["out"],
+)
+def fn_with_numpy_operation(x: Tensor, y: Tensor) -> Tensor:
+    x_np = x.numpy()
+    y_np = y.numpy()
+    x_np_reduce = x_np.sum(axis=-1)
+    y_np_reduce = y_np.sum(axis=-1)
+    return paddle.to_tensor(x_np_reduce + y_np_reduce).cast(paddle.int32)
+
+
+@paddle.static.register_op(
+    name="fn_with_constant",
+    infer_meta=lambda x, c: paddle.static.MetaTensor(
+        dtype=x.dtype, shape=[np.prod(x.shape).item()]
+    ),
+    input_names=["x"],
+    output_names=["out"],
+)
+def fn_with_constant(x: paddle.Tensor, c: int):
+    return paddle.to_tensor(x.flatten() + c)
+
+
+class PythonOpTestMixin:
+    inputs: dict[str, paddle.Tensor]
+    constants: dict[str, object]
+    fn: Callable[..., paddle.Tensor]
+
+    def run_in_dygraph(self):
+        return self.fn(**self.inputs, **self.constants)
+
+    @static_guard()
+    def run_in_static(self):
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
+            input_values = {
+                k: paddle.static.data(name=k, shape=v.shape, dtype=v.dtype)
+                for k, v in self.inputs.items()
+            }
+            out_value = self.fn(**input_values, **self.constants)
+            exe = paddle.static.Executor()
+        (out,) = exe.run(
+            main_program,
+            feed={k: v.numpy() for k, v in self.inputs.items()},
+            fetch_list=[out_value],
+        )
+        return out
+
+    def test_dy_st(self):
+        np.testing.assert_allclose(self.run_in_dygraph(), self.run_in_static())
+
+
+class TestFnWithBreakgraph(unittest.TestCase, PythonOpTestMixin):
+    def setUp(self):
+        self.fn = fn_with_breakgraph
+        self.inputs = {
+            "x": paddle.randn([2, 3, 4]),
+            "y": paddle.randn([2, 3, 4]),
+        }
+        self.constants = {}
+
+
+class TestFnWithNumPyOperation(unittest.TestCase, PythonOpTestMixin):
+    def setUp(self):
+        self.fn = fn_with_numpy_operation
+        self.inputs = {
+            "x": paddle.randn([7, 8, 9]),
+            "y": paddle.randn([7, 8, 9]),
+        }
+        self.constants = {}
+
+
+class TestFnWithConstant1(unittest.TestCase, PythonOpTestMixin):
+    def setUp(self):
+        self.fn = fn_with_constant
+        self.inputs = {
+            "x": paddle.randn([4, 5, 6]),
+        }
+        self.constants = {"c": -1}
+
+
+class TestFnWithConstant2(unittest.TestCase, PythonOpTestMixin):
+    def setUp(self):
+        self.fn = fn_with_constant
+        self.inputs = {
+            "x": paddle.randn([4, 5, 6]),
+        }
+        self.constants = {"c": -2}  # Note that hash(-1) == hash(-2)
+
+
+def fn_use_2_register_op(x: Tensor, y: Tensor) -> Tensor:
+    z1 = fn_with_breakgraph(x, y)
+    z2 = fn_with_numpy_operation(x, y)
+    out = z1 * 100 + z2.unsqueeze(axis=-1).astype(paddle.float32)
+    return out
+
+
+class TestToStatic(Dy2StTestBase):
+    def test_to_static_use_2_op(self):
+        x = paddle.randn([2, 3, 4])
+        y = paddle.randn([2, 3, 4])
+        fn = fn_use_2_register_op
+        dy_out = fn(x, y)
+        static_fn = paddle.jit.to_static(fn)
+        st_out = static_fn(x, y)
+        np.testing.assert_allclose(dy_out, st_out)
+
+
+if __name__ == "__main__":
+    unittest.main()