Могу ли я использовать pybind11 для передачи пустого массива в функцию, принимающую Eigen :: Tensor?

Question

Могу ли я использовать pybind1 для передачи трехмерного массива numpy в функцию c ++, принимающую Eigen::Tensor в качестве аргумента. Например, рассмотрим следующую функцию c ++:

Eigen::Tensor<double, 3> addition_tensor(Eigen::Tensor<double, 3> a,
                                         Eigen::Tensor<double, 3> b) {
    return a + b;
}

после компиляции функции, импорта ее в python и передачи ей массива numpy np.ones((1, 2, 2)) я получаю следующее сообщение об ошибке:

TypeError: addition_tensor(): incompatible function arguments. The following argument types are supported:
    1. (arg0: Eigen::Tensor<double, 3, 0, long>, arg1: Eigen::Tensor<double, 3, 0, long>) -> Eigen::Tensor<double, 3, 0, long>

Я особенно удивлен тем, что не могу передать трехмерный массив NumPy, поскольку я могу передать двумерный numpy array функции, принимающей Eigen::MatrixXd, как:

Eigen::MatrixXd addition(Eigen::MatrixXd a, Eigen::MatrixXd b) { return a + b; }

Весь код, который я использовал для этого примера:

#include <eigen-git-mirror/Eigen/Dense>
#include <eigen-git-mirror/unsupported/Eigen/CXX11/Tensor>
#include "pybind11/include/pybind11/eigen.h"
#include "pybind11/include/pybind11/pybind11.h"

Eigen::MatrixXd addition(Eigen::MatrixXd a, Eigen::MatrixXd b) { return a + b; }

Eigen::Tensor<double, 3> addition_tensor(Eigen::Tensor<double, 3> a,
                                         Eigen::Tensor<double, 3> b) {
    return a + b;
}

PYBIND11_MODULE(example, m) {
    m.def("addition", &addition, "A function which adds two numbers");
    m.def("addition_tensor", &addition_tensor,
          "A function which adds two numbers");
}

Я скомпилировал приведенный выше код с помощью g++ -shared -fPIC `python3 -m pybind11 --includes` example.cpp -o example`python3-config --extension-suffix`. У кого-нибудь есть идея, как я могу использовать трехмерный массив numpy для функции, принимающей трехмерный Eigen::Tensor?

Answer 1

Это не поддерживается напрямую, вот некоторые обсуждения (включая некоторый код для сопоставления, если вы хотите добавить это в свой проект): https://github.com/pybind/pybind11/issues/1377

Answer 2

Спасибо за ответ @Джона Цвинка, я смог достичь того, что искал. Если кому-то интересно, вот репликация:

#include <eigen-git-mirror/Eigen/Dense>
#include <eigen-git-mirror/unsupported/Eigen/CXX11/Tensor>
#include "pybind11/include/pybind11/eigen.h"
#include "pybind11/include/pybind11/numpy.h"
#include "pybind11/include/pybind11/pybind11.h"

Eigen::Tensor<double, 3, Eigen::RowMajor> getTensor(
    pybind11::array_t<double> inArray) {
    // request a buffer descriptor from Python
    pybind11::buffer_info buffer_info = inArray.request();

    // extract data an shape of input array
    double *data = static_cast<double *>(buffer_info.ptr);
    std::vector<ssize_t> shape = buffer_info.shape;

    // wrap ndarray in Eigen::Map:
    // the second template argument is the rank of the tensor and has to be
    // known at compile time
    Eigen::TensorMap<Eigen::Tensor<double, 3, Eigen::RowMajor>> in_tensor(
        data, shape[0], shape[1], shape[2]);
    return in_tensor;
}

pybind11::array_t<double> return_array(
    Eigen::Tensor<double, 3, Eigen::RowMajor> inp) {
    std::vector<ssize_t> shape(3);
    shape[0] = inp.dimension(0);
    shape[1] = inp.dimension(1);
    shape[2] = inp.dimension(2);
    return pybind11::array_t<double>(
        shape,  // shape
        {shape[1] * shape[2] * sizeof(double), shape[2] * sizeof(double),
         sizeof(double)},  // strides
        inp.data());       // data pointer
}

pybind11::array_t<double> addition(pybind11::array_t<double> a,
                                   pybind11::array_t<double> b) {
    Eigen::Tensor<double, 3, Eigen::RowMajor> a_t = getTensor(a);
    Eigen::Tensor<double, 3, Eigen::RowMajor> b_t = getTensor(b);
    Eigen::Tensor<double, 3, Eigen::RowMajor> res = a_t + b_t;
    return return_array(res);
}

PYBIND11_MODULE(example, m) {
    m.def("addition", &addition, "A function which adds two numbers");
}

В отличие от предложения в ссылке, на которую ссылался Джон, я не возражал против использования RowMajor порядка хранения для Eigen::Tensor. Я видел, что этот порядок хранения использовался несколько раз и в коде tensorflow. Я не знаю, копирует ли приведенный выше код ненужные данные.