python - py::vectorize + type_caster = NumPy 类型信息缺失

Question

在过去的几天里，我一直在使用pybind11为现有的 C++ 库创建 Python 绑定(bind)，我真的很喜欢它!

不幸的是，我刚刚遇到了一个小问题......

我想做两件事:

• 自定义 type_caster，用于将第三方 vector 类型转换为 NumPy 数组并返回

• 返回此类型的函数，由 py::vectorize() 自动矢量化

这两件事本身都能很好地工作。具有标量输入的矢量化函数也能很好地工作。

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

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

我做错了什么？

或者这根本不起作用？

<小时/>

以下是我精简到最少的代码。在我的实际代码中，vec 类是第三方库的一部分，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() 还不支持返回 np::array 的函数。

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