From 9796fe98fca282a2abe3719057397e256ba2e290 Mon Sep 17 00:00:00 2001 From: Yannick Jadoul Date: Fri, 2 Oct 2020 21:30:34 +0200 Subject: [PATCH] feat: vectorize functions with void return type (#1969) * Allow function/functor passed to py::vectorize to return void * Stealing @sizmailov's test and fixing unused argument warning * Add missing std::move() RVO doesn't work here because function return type is different from actual returned type * remove extra EOL * docs: add a few details * chore: pre-commit autoupdate * Remove array_iterator, array_begin, and array_end (in detail namespace) Co-authored-by: Sergei Izmailov Co-authored-by: Henry Schreiner --- docs/changelog.rst | 17 ++++--- include/pybind11/numpy.h | 91 ++++++++++++++++++++++++---------- tests/test_numpy_vectorize.cpp | 6 ++- tests/test_numpy_vectorize.py | 11 ++++ 4 files changed, 90 insertions(+), 35 deletions(-) diff --git a/docs/changelog.rst b/docs/changelog.rst index 56dc10eaf..b19cd2eab 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -98,6 +98,9 @@ See :ref:`upgrade-guide-2.6` for help upgrading to the new version. ``get_type_overload`` is deprecated. `#2325 `_ +* Error now thrown when ``__init__`` is forgotten on subclasses. + `#2152 `_ + * `py::class_` is now supported. Note that writing to one data member of the union and reading another (type punning) is UB in C++. Thus pybind11-bound enums should never be used for such conversion. @@ -109,9 +112,6 @@ Smaller or developer focused features: .. _pybind11-mkdoc: https://github.com/pybind/pybind11-mkdoc -* Error now thrown when ``__init__`` is forgotten on subclasses. - `#2152 `_ - * If ``__eq__`` defined but not ``__hash__``, ``__hash__`` is now set to ``None``. `#2291 `_ @@ -122,9 +122,6 @@ Smaller or developer focused features: * Throw if conversion to ``str`` fails. `#2477 `_ -* Added missing signature for ``py::array``. - `#2363 `_ - * Pointer to ``std::tuple`` & ``std::pair`` supported in cast. `#2334 `_ @@ -132,7 +129,13 @@ Smaller or developer focused features: argument type. `#2293 `_ -* Bugfixes related to more extensive testing +* Added missing signature for ``py::array``. + `#2363 `_ + +* ``py::vectorize`` is now supported on functions that return void. + `#1969 `_ + +* Bugfixes related to more extensive testing. `#2321 `_ * Bug in timezone issue in Eastern hemisphere midnight fixed. diff --git a/include/pybind11/numpy.h b/include/pybind11/numpy.h index ceb06fe36..54e160f35 100644 --- a/include/pybind11/numpy.h +++ b/include/pybind11/numpy.h @@ -1274,19 +1274,6 @@ private: #endif // __CLION_IDE__ -template -using array_iterator = typename std::add_pointer::type; - -template -array_iterator array_begin(const buffer_info& buffer) { - return array_iterator(reinterpret_cast(buffer.ptr)); -} - -template -array_iterator array_end(const buffer_info& buffer) { - return array_iterator(reinterpret_cast(buffer.ptr) + buffer.size); -} - class common_iterator { public: using container_type = std::vector; @@ -1486,6 +1473,56 @@ struct vectorize_arg { using type = conditional_t, array::forcecast>, T>; }; + +// py::vectorize when a return type is present +template +struct vectorize_returned_array { + using Type = array_t; + + static Type create(broadcast_trivial trivial, const std::vector &shape) { + if (trivial == broadcast_trivial::f_trivial) + return array_t(shape); + else + return array_t(shape); + } + + static Return *mutable_data(Type &array) { + return array.mutable_data(); + } + + static Return call(Func &f, Args &... args) { + return f(args...); + } + + static void call(Return *out, size_t i, Func &f, Args &... args) { + out[i] = f(args...); + } +}; + +// py::vectorize when a return type is not present +template +struct vectorize_returned_array { + using Type = none; + + static Type create(broadcast_trivial, const std::vector &) { + return none(); + } + + static void *mutable_data(Type &) { + return nullptr; + } + + static detail::void_type call(Func &f, Args &... args) { + f(args...); + return {}; + } + + static void call(void *, size_t, Func &f, Args &... args) { + f(args...); + } +}; + + template struct vectorize_helper { @@ -1520,6 +1557,8 @@ private: using arg_call_types = std::tuple::call_type...>; template using param_n_t = typename std::tuple_element::type; + using returned_array = vectorize_returned_array; + // Runs a vectorized function given arguments tuple and three index sequences: // - Index is the full set of 0 ... (N-1) argument indices; // - VIndex is the subset of argument indices with vectorized parameters, letting us access @@ -1551,20 +1590,19 @@ private: // not wrapped in an array). if (size == 1 && ndim == 0) { PYBIND11_EXPAND_SIDE_EFFECTS(params[VIndex] = buffers[BIndex].ptr); - return cast(f(*reinterpret_cast *>(params[Index])...)); + return cast(returned_array::call(f, *reinterpret_cast *>(params[Index])...)); } - array_t result; - if (trivial == broadcast_trivial::f_trivial) result = array_t(shape); - else result = array_t(shape); + auto result = returned_array::create(trivial, shape); if (size == 0) return std::move(result); /* Call the function */ + auto mutable_data = returned_array::mutable_data(result); if (trivial == broadcast_trivial::non_trivial) - apply_broadcast(buffers, params, result, i_seq, vi_seq, bi_seq); + apply_broadcast(buffers, params, mutable_data, size, shape, i_seq, vi_seq, bi_seq); else - apply_trivial(buffers, params, result.mutable_data(), size, i_seq, vi_seq, bi_seq); + apply_trivial(buffers, params, mutable_data, size, i_seq, vi_seq, bi_seq); return std::move(result); } @@ -1587,7 +1625,7 @@ private: }}; for (size_t i = 0; i < size; ++i) { - out[i] = f(*reinterpret_cast *>(params[Index])...); + returned_array::call(out, i, f, *reinterpret_cast *>(params[Index])...); for (auto &x : vecparams) x.first += x.second; } } @@ -1595,19 +1633,18 @@ private: template void apply_broadcast(std::array &buffers, std::array ¶ms, - array_t &output_array, + Return *out, + size_t size, + const std::vector &output_shape, index_sequence, index_sequence, index_sequence) { - buffer_info output = output_array.request(); - multi_array_iterator input_iter(buffers, output.shape); + multi_array_iterator input_iter(buffers, output_shape); - for (array_iterator iter = array_begin(output), end = array_end(output); - iter != end; - ++iter, ++input_iter) { + for (size_t i = 0; i < size; ++i, ++input_iter) { PYBIND11_EXPAND_SIDE_EFFECTS(( params[VIndex] = input_iter.template data() )); - *iter = f(*reinterpret_cast *>(std::get(params))...); + returned_array::call(out, i, f, *reinterpret_cast *>(std::get(params))...); } } }; diff --git a/tests/test_numpy_vectorize.cpp b/tests/test_numpy_vectorize.cpp index e76e462cb..1f2de5fd2 100644 --- a/tests/test_numpy_vectorize.cpp +++ b/tests/test_numpy_vectorize.cpp @@ -50,7 +50,9 @@ TEST_SUBMODULE(numpy_vectorize, m) { NonPODClass(int v) : value{v} {} int value; }; - py::class_(m, "NonPODClass").def(py::init()); + py::class_(m, "NonPODClass") + .def(py::init()) + .def_readwrite("value", &NonPODClass::value); m.def("vec_passthrough", py::vectorize( [](double *a, double b, py::array_t c, const int &d, int &e, NonPODClass f, const double g) { return *a + b + c.at(0) + d + e + f.value + g; @@ -86,4 +88,6 @@ TEST_SUBMODULE(numpy_vectorize, m) { std::array buffers {{ arg1.request(), arg2.request(), arg3.request() }}; return py::detail::broadcast(buffers, ndim, shape); }); + + m.def("add_to", py::vectorize([](NonPODClass& x, int a) { x.value += a; })); } diff --git a/tests/test_numpy_vectorize.py b/tests/test_numpy_vectorize.py index 54e44cd8d..992ff8ecc 100644 --- a/tests/test_numpy_vectorize.py +++ b/tests/test_numpy_vectorize.py @@ -192,3 +192,14 @@ def test_array_collapse(): z = m.vectorized_func(1, [[[2]]], 3) assert isinstance(z, np.ndarray) assert z.shape == (1, 1, 1) + + +def test_vectorized_noreturn(): + x = m.NonPODClass(0) + assert x.value == 0 + m.add_to(x, [1, 2, 3, 4]) + assert x.value == 10 + m.add_to(x, 1) + assert x.value == 11 + m.add_to(x, [[1, 1], [2, 3]]) + assert x.value == 18