[CppExtension Unittest] Add unit test of vector<Tensor> (#53040)

PaddlePaddle · Apr 24, 2023 · b06ec0c · b06ec0c
1 parent 81fb7df
commit b06ec0c
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 0 deletions.
diff --git 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<paddle::Tensor> custom_tensor(
+    const std::vector<paddle::Tensor>& inputs) {
+  std::vector<paddle::Tensor> 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<paddle::Tensor> 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
@@ -68,6 +68,7 @@ def tearDown(self):
     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 @@ def _test_extension_class(self):
                 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
@@ -148,6 +148,7 @@ def test_cpp_extension(self):
         # 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 @@ def _test_extension_class(self):
                 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