diff --git a/python/paddle/fluid/tests/cpp_extension/custom_extension.cc b/python/paddle/fluid/tests/cpp_extension/custom_extension.cc index 2fc5c42a80d75c0f4571cca8e0b492ef0f1b8d25..67653bffb3bdb31931ae46d813d5514830c4609f 100644 --- a/python/paddle/fluid/tests/cpp_extension/custom_extension.cc +++ b/python/paddle/fluid/tests/cpp_extension/custom_extension.cc @@ -26,6 +26,16 @@ paddle::Tensor custom_add(const paddle::Tensor& x, const paddle::Tensor& y) { return x.exp() + y.exp(); } +std::vector custom_tensor( + const std::vector& inputs) { + std::vector out; + out.reserve(inputs.size()); + for (const auto& input : inputs) { + out.push_back(input + 1.0); + } + return out; +} + paddle::Tensor nullable_tensor(bool return_none = false) { paddle::Tensor t; if (!return_none) { @@ -45,6 +55,7 @@ paddle::optional optional_tensor(bool return_option = false) { PYBIND11_MODULE(custom_cpp_extension, m) { m.def("custom_add", &custom_add, "exp(x) + exp(y)"); m.def("custom_sub", &custom_sub, "exp(x) - exp(y)"); + m.def("custom_tensor", &custom_tensor, "x + 1"); m.def("nullable_tensor", &nullable_tensor, "returned Tensor might be None"); m.def( "optional_tensor", &optional_tensor, "returned Tensor might be optional"); diff --git a/python/paddle/fluid/tests/cpp_extension/test_cpp_extension_jit.py b/python/paddle/fluid/tests/cpp_extension/test_cpp_extension_jit.py index bc6f8113afd9129a1e4ebd372dd03e6d9ea2a1c7..5723df1b585b0428f248ec613b3baec6d726da14 100644 --- a/python/paddle/fluid/tests/cpp_extension/test_cpp_extension_jit.py +++ b/python/paddle/fluid/tests/cpp_extension/test_cpp_extension_jit.py @@ -68,6 +68,7 @@ class TestCppExtensionJITInstall(unittest.TestCase): def test_cpp_extension(self): self._test_extension_function() self._test_extension_class() + self._test_vector_tensor() self._test_nullable_tensor() self._test_optional_tensor() if paddle.is_compiled_with_cuda(): @@ -109,6 +110,20 @@ class TestCppExtensionJITInstall(unittest.TestCase): atol=1e-5, ) + def _test_vector_tensor(self): + for dtype in self.dtypes: + np_inputs = [ + np.random.uniform(-1, 1, [4, 8]).astype(dtype) for _ in range(3) + ] + inputs = [paddle.to_tensor(np_x, dtype=dtype) for np_x in np_inputs] + + out = custom_cpp_extension.custom_tensor(inputs) + target_out = [x + 1.0 for x in inputs] + for i in range(3): + np.testing.assert_allclose( + out[i].numpy(), target_out[i].numpy(), atol=1e-5 + ) + def _test_nullable_tensor(self): x = custom_cpp_extension.nullable_tensor(True) assert x is None, "Return None when input parameter return_none = True" diff --git a/python/paddle/fluid/tests/cpp_extension/test_cpp_extension_setup.py b/python/paddle/fluid/tests/cpp_extension/test_cpp_extension_setup.py index 53dffde43209599c7275ed78d26c5ddd4d6ee6ef..15d0cb77d03d5cc42fad61d223badf1a118a4627 100644 --- a/python/paddle/fluid/tests/cpp_extension/test_cpp_extension_setup.py +++ b/python/paddle/fluid/tests/cpp_extension/test_cpp_extension_setup.py @@ -148,6 +148,7 @@ class TestCppExtensionSetupInstall(unittest.TestCase): # Extension self._test_extension_function_plain() self._test_extension_function_mixed() + self._test_vector_tensor() self._test_extension_class() self._test_nullable_tensor() self._test_optional_tensor() @@ -218,6 +219,22 @@ class TestCppExtensionSetupInstall(unittest.TestCase): atol=1e-5, ) + def _test_vector_tensor(self): + import custom_cpp_extension + + for dtype in self.dtypes: + np_inputs = [ + np.random.uniform(-1, 1, [4, 8]).astype(dtype) for _ in range(3) + ] + inputs = [paddle.to_tensor(np_x, dtype=dtype) for np_x in np_inputs] + + out = custom_cpp_extension.custom_tensor(inputs) + target_out = [x + 1 for x in inputs] + for i in range(3): + np.testing.assert_allclose( + out[i].numpy(), target_out[i].numpy(), atol=1e-5 + ) + def _test_nullable_tensor(self): import custom_cpp_extension