From b06ec0c759cf79aa92d6124ac6fcde315785574f Mon Sep 17 00:00:00 2001
From: HongyuJia <jiahongyu@baidu.com>
Date: Mon, 24 Apr 2023 16:02:21 +0800
Subject: [PATCH] [CppExtension Unittest] Add unit test of vector<Tensor>
 (#53040)

---
 .../tests/cpp_extension/custom_extension.cc     | 11 +++++++++++
 .../cpp_extension/test_cpp_extension_jit.py     | 15 +++++++++++++++
 .../cpp_extension/test_cpp_extension_setup.py   | 17 +++++++++++++++++
 3 files changed, 43 insertions(+)

diff --git a/python/paddle/fluid/tests/cpp_extension/custom_extension.cc b/python/paddle/fluid/tests/cpp_extension/custom_extension.cc
index 2fc5c42a80d..67653bffb3b 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<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
index bc6f8113afd..5723df1b585 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 53dffde4320..15d0cb77d03 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
 
-- 
GitLab