py::vectorize + type_caster = NumPy type info missing

1.2k views Asked by At

In the last few days, I've been using pybind11 to create Python bindings for an existing C++ library, and I really like it!

Sadly, I've just run into a little problem ...

I'm trying to have two things:

  • A custom type_caster that converts a third-party vector type to NumPy arrays and back

  • A function returning this type, which is automatically vectorized by py::vectorize()

Both things on their own work nicely. The vectorized function with scalar input also works nicely.

However, if I call the vectorized function with an array as input, an exception is raised:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: NumPy type info missing for 3vecIdLi2EE

What am I doing wrong?

Or isn't this supposed to work at all?


The following is my code reduced to a minimum. In my actual code, the vec class is part of a third party library and return_vector() is in my own code.

mylib.cpp:

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>

namespace py = pybind11;

template<typename T, int N> struct vec {
  explicit vec(const T* data_) {
    for (int i = 0; i < N; ++i) { this->data[i] = data_[i]; }
  }
  T data[N];
};

vec<double, 2> return_vector(double t) {
  double v[] = {t, t};
  return vec<double, 2>{v};
}

namespace pybind11 { namespace detail {
  template <typename T, int N> struct type_caster<vec<T, N>>
  {
  private:
    using _vecTN = vec<T, N>;

  public:
    PYBIND11_TYPE_CASTER(_vecTN, _("vec<T, N>"));

    bool load(py::handle src, bool convert)
    {
      if (!convert && !py::array_t<T>::check_(src)) { return false; }
      auto buf = py::array_t<T>::ensure(src);
      if (!buf || buf.ndim() != 1 || buf.size() != N) { return false; }
      value = _vecTN{buf.data()};
      return true;
    }

    static py::handle cast(const _vecTN& src,
        py::return_value_policy policy, py::handle parent)
    {
      py::array_t<T> a({N});
      for (auto i = 0; i < N; ++i) { a.mutable_at(i) = src.data[i]; }
      return a.release();
    }
  };
}}

template struct pybind11::detail::type_caster<vec<double, 2>>;

PYBIND11_MODULE(mylib, m) {
  m.def("return_vector", py::vectorize(&return_vector));
}

(Feel free to comment on the code, I might be doing many things wrong. I'm especially unsure about my type_caster code.)

For completeness, here's the corresponding setup.py:

from setuptools import setup, Extension

class get_pybind_include(object):

    def __init__(self, user=False):
        self.user = user

    def __str__(self):
        import pybind11
        return pybind11.get_include(self.user)

ext_modules = [
    Extension(
        'mylib',
        ['mylib.cpp'],
        include_dirs=[
            get_pybind_include(),
            get_pybind_include(user=True),
        ],
        language='c++',
    ),
]

setup(
    name='mylib',
    ext_modules=ext_modules,
    install_requires=['pybind11>=2.2'],
)

I've compiled the extension module with

python3 setup.py develop

Running this Python code works fine:

>>> import mylib
>>> mylib.return_vector(1)
array([1., 1.])

However, when I call it with an array input, I get an error:

>>> mylib.return_vector([2, 3])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: NumPy type info missing for 3vecIdLi2EE

I would have hoped for a 2-dimensional array, something like:

array([[2., 2.],
       [3., 3.]])
1

There are 1 answers

0
Matthias On BEST ANSWER

It turns out that py::vectorize() doesn't (yet?) support functions that return a np::array.

See https://github.com/pybind/pybind11/issues/763.