py::vectorize + type_caster = 缺少 NumPy 类型信息

Question

最近几天，我一直在使用 pybind11 为现有的 C++ 库创建 Python 绑定，我真的很喜欢它！

遗憾的是，我刚刚运行遇到了一个小问题...

我正在尝试做两件事：

将第三方矢量类型转换为 NumPy 数组并返回的自定义 type_caster
返回该类型的函数，由py::vectorize()

这两件事各自都能很好地工作。带有标量输入的向量化函数也能很好地工作。

但是，如果我使用数组作为输入调用向量化函数，则会引发异常：

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

我做错了什么？

或者这根本不起作用？

以下是我精简的代码。在我的实际代码中，vec class 是第三方库的一部分，而 return_vector() 在我自己的代码中。

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));
}

（欢迎对代码发表评论，我可能做错了很多事情。我对我的 type_caster 代码特别不确定。）

为了完整起见，这里是相应的 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'],
)

我用

编译了扩展模块

python3 setup.py develop

运行这个 Python 代码工作正常：

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

但是，当我用数组输入调用它时，出现错误：

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

我希望有一个二维数组，类似于：

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

Answer 1

原来 py::vectorize() 不（还？）支持 return 和 np::array 的功能。

参见https://github.com/pybind/pybind11/issues/763。

py::vectorize + type_caster = 缺少 NumPy 类型信息

py::vectorize + type_caster = NumPy type info missing

c++

python

pybind11