[Cpp Extension] Support optional types (#50764)

* [Cpp Extension] Support optional type

* fix custom_extension.cc

Author: jiahy0825

paddle/utils/pybind.h

@@ -15,6 +15,7 @@
 #pragma once
 
 #include "paddle/phi/api/include/tensor.h"
+#include "paddle/utils/optional.h"
 #include "pybind11/pybind11.h"
 #include "pybind11/stl.h"
 
@@ -73,5 +74,12 @@ struct type_caster<paddle::experimental::Tensor> {
         src, true /* return_py_none_if_not_initialize */));
   }
 };
+
+// Pybind11 bindings for optional types.
+// http://pybind11.readthedocs.io/en/stable/advanced/cast/stl.html#c-17-library-containers
+template <typename T>
+struct type_caster<paddle::optional<T>> : optional_caster<paddle::optional<T>> {
+};
+
 }  // namespace detail
 }  // namespace pybind11

python/paddle/fluid/tests/cpp_extension/custom_extension.cc

@@ -32,10 +32,20 @@ paddle::Tensor nullable_tensor(bool return_none = false) {
   return t;
 }
 
+paddle::optional<paddle::Tensor> optional_tensor(bool return_option = false) {
+  paddle::optional<paddle::Tensor> t;
+  if (!return_option) {
+    t = paddle::ones({2, 2});
+  }
+  return t;
+}
+
 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("nullable_tensor", &nullable_tensor, "returned Tensor might be None");
+  m.def(
+      "optional_tensor", &optional_tensor, "returned Tensor might be optional");
 
   py::class_<Power>(m, "Power")
       .def(py::init<int, int>())

python/paddle/fluid/tests/cpp_extension/test_cpp_extension_jit.py

@@ -67,6 +67,8 @@ def tearDown(self):
     def test_cpp_extension(self):
         self._test_extension_function()
         self._test_extension_class()
+        self._test_nullable_tensor()
+        self._test_optional_tensor()
 
     def _test_extension_function(self):
         for dtype in self.dtypes:
@@ -104,6 +106,30 @@ def _test_extension_class(self):
                 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"
+        x = custom_cpp_extension.nullable_tensor(False).numpy()
+        x_np = np.ones(shape=[2, 2])
+        np.testing.assert_array_equal(
+            x,
+            x_np,
+            err_msg='extension out: {},\n numpy out: {}'.format(x, x_np),
+        )
+
+    def _test_optional_tensor(self):
+        x = custom_cpp_extension.optional_tensor(True)
+        assert (
+            x is None
+        ), "Return None when input parameter return_option = True"
+        x = custom_cpp_extension.optional_tensor(False).numpy()
+        x_np = np.ones(shape=[2, 2])
+        np.testing.assert_array_equal(
+            x,
+            x_np,
+            err_msg='extension out: {},\n numpy out: {}'.format(x, x_np),
+        )
+
 
 if __name__ == '__main__':
     unittest.main()

python/paddle/fluid/tests/cpp_extension/test_cpp_extension_setup.py

@@ -149,6 +149,7 @@ def test_cpp_extension(self):
         self._test_extension_function_mixed()
         self._test_extension_class()
         self._test_nullable_tensor()
+        self._test_optional_tensor()
         # Custom op
         self._test_static()
         self._test_dynamic()
@@ -227,6 +228,21 @@ def _test_nullable_tensor(self):
             err_msg='extension out: {},\n numpy out: {}'.format(x, x_np),
         )
 
+    def _test_optional_tensor(self):
+        import custom_cpp_extension
+
+        x = custom_cpp_extension.optional_tensor(True)
+        assert (
+            x is None
+        ), "Return None when input parameter return_option = True"
+        x = custom_cpp_extension.optional_tensor(False).numpy()
+        x_np = np.ones(shape=[2, 2])
+        np.testing.assert_array_equal(
+            x,
+            x_np,
+            err_msg='extension out: {},\n numpy out: {}'.format(x, x_np),
+        )
+
     def _test_static(self):
         import mix_relu_extension