wehub-resource-sync
160 lines
4.5 KiB
Python
160 lines
4.5 KiB
Python
# 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()
|