frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add optional support to C++ extensions (#7055)

Browse Source
This commit is contained in:
Peter Goldsborough
2018-04-28 01:59:50 +01:00
committed by GitHub
parent 7b09bc72a5
commit 2e023a29e4
3 changed files with 17 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
test/cpp_extensions/extension.cpp
View File

@ -20,8 +20,16 @@ struct MatrixMultiplier {
at::Tensor tensor_;
};
bool function_taking_optional(at::optional<at::Tensor> tensor) {
return tensor.has_value();
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("sigmoid_add", &sigmoid_add, "sigmoid(x) + sigmoid(y)");
m.def(
"function_taking_optional",
&function_taking_optional,
"function_taking_optional");
py::class_<MatrixMultiplier>(m, "MatrixMultiplier")
.def(py::init<int, int>())
.def("forward", &MatrixMultiplier::forward)

6
test/test_cpp_extensions.py
View File

@ -100,6 +100,12 @@ class TestCppExtension(common.TestCase):
# 2 * sigmoid(0) = 2 * 0.5 = 1
self.assertEqual(z, torch.ones_like(z))
def test_optional(self):
has_value = cpp_extension.function_taking_optional(torch.ones(5))
self.assertTrue(has_value)
has_value = cpp_extension.function_taking_optional(None)
self.assertFalse(has_value)
if __name__ == '__main__':
common.run_tests()

3
torch/csrc/utils/pybind.h
View File

@ -57,4 +57,7 @@ public:
}
};
// http://pybind11.readthedocs.io/en/stable/advanced/cast/stl.html#c-17-library-containers
template <typename T>
struct type_caster<at::optional<T>> : optional_caster<at::optional<T>> {};
}} // namespace pybind11::detail