Can you support Tensor datatype from Eigen library?

[Felix-neko]

Hi folks! And thank you for your brilliant library.

I need to pass 3-dimensional arrays to my pybind-made Python module that wraps some C++ code. I want to pass it a 3D numpy.ndarray and automatically convert it into Eigen Tensor.

Can you support it?

[psiha]

Please also add support for Eigen::TensorMap at the same time.

I added support for these in my code 'the pybind::class_< Eigen::Tensor... > way' and it works but not optimally:

• there is no implicit conversion to-from numpy ndarrays (conversion works but I has to be done explicitly on Python's side) like there is with the builtin (conversion) support for Eigen::Matrix types
• there is no automatic copy elision (np.array( x, copy = False ))
• I couldn't find a way to specify/pass constness through the buffer_info interface...
  Could you maybe provide a third, intermediate way of defining type conversions/mappings - that's more powerful/lower level than class_ but also simpler/higher level than the C Python API level as described in the 'Custom type casters' section of the documentation, that would allow for the above three points to be resolved?

[jagerman]

Could you maybe provide a third, intermediate way of defining type conversions/mappings

That's basically what the whole type_caster specialization system is for. See how the existing implementations in eigen.h work; that's where additional specializations for Tensor would go.

[JonasBreuling]

I'm not sure if this is still important to you, but its pretty easy to wrap an input ndarray using Eigen::TensorMap<Eigen::Tensor<T, dims...>> and pybind11.

Here comes the cpp file which returns and rank 3 ndarray with reversed ordering of an input ndarray of rank 3:

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <unsupported/Eigen/CXX11/Tensor>

namespace py = pybind11;

// return 3 dimensional ndarray with reversed order of input ndarray
template<class T>
py::array_t<T> eigenTensor(py::array_t<T> inArray) {

    // request a buffer descriptor from Python
    py::buffer_info buffer_info = inArray.request();

    // extract data an shape of input array
    T *data = static_cast<T *>(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<T, 3>> in_tensor(data, shape[0], shape[1], shape[2]);

    // build result tensor with reverse ordering
    Eigen::Tensor<T, 3> out_tensor(2, 2, 2);
    for (int i=0; i < shape[0]; i++) {
        for (int j=0; j < shape[1]; j++) {
            for (int k=0; k < shape[2]; k++) {
                out_tensor(k, j, i) = in_tensor(i, j, k);
            }
        }

    }

    // return numpy array wrapping eigen tensor's pointer
    return py::array_t<T>(shape, // shape
                          {shape[0] * shape[1] * sizeof(T), shape[1] * sizeof(T), sizeof(T)},  // strides
                          out_tensor.data());  // data pointer
}

PYBIND11_MODULE(eigen, m) {
    m.def("eigenTensor", &eigenTensor<double>, py::return_value_policy::move,
        py::arg("inArray"));
}

And you can test it with

import numpy as np
from eigen import eigenTensor

def eigenTensorPython(inArray):
    outArray = np.zeros((2, 2, 2))
    for i in range(inArray.shape[0]):
        for j in range(inArray.shape[0]):
            for k in range(inArray.shape[0]):
                outArray[k, j, i] = inArray[i, j, k]
    return outArray

n = int(1.0e5)

# initialize unique ndarray
inArray = np.zeros((2, 2, 2))
for i in range(inArray.shape[0]):
    for j in range(inArray.shape[0]):
        for k in range(inArray.shape[0]):
            inArray[i, j, k] = i * 2 * 2 + j * 2 + k

print(inArray)

# print(eigenTensor(inArray))
# print(eigenTensorPython(inArray))
assert( (eigenTensorPython(inArray) == eigenTensor(inArray)).all() )

If it's not of interest for you it might help others to use Eigen::Tensor and numpy ndarray together with pybind11. I think these conversations are that easy so you don't need it to be automated. For speed reasons you might think of Eigen::TensorFixedSize , see here.

Nota bene: As written in the doc Eigen::Tensor<T, dims...> supports 2 layouts: ColMajor (the default) and RowMajor. Only the default column major layout is currently fully supported, and it is therefore not recommended to attempt to use the row major layout at the moment. But data in an ndarray is stored in the row-major (C) order, see doc. This implies that you have to keept that in mind. Either you have to use Eigen::Tensor<T, dims..., Eigen::RowMajor> which is not fully supported on the C++ side, or you have to return and ndarray with Fortran-style contiguous arrays is set to true.

Best wishes
Jonas

[ghost]

@JonasHarsch , so is possible to pass a Eigen::Tensor from C++ to Python...and get back from Python to C++ using Pybind11?

[JonasBreuling]

Yes, at least it was when I wrote the answer above. Maybe it's best if you try to get my example code working. Tell me if you have problems, then I can try it myself tomorrow.

[bhaveshshrimali]

Hi everyone,

I just started working with pybind11 and was trying to experiment a bit with calculating the determinant of a stack of NumPy arrays, say an array of shape (3,3,100,4) with the determinant broadcasted over the trailing axis so the result is an array of shape (100,4). This was of course in hopes of getting some speed-up.

I was trying to tweak the above example by @JonasHarsch with something like:

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <unsupported/Eigen/CXX11/Tensor>

namespace py = pybind11;

// return 4 dimensional ndarray with determinant broadcasted across trailing axis
template<class T>
py::array_t<T> calculateDeterminant(py::array_t<T> inArray) {

    // request a buffer descriptor from Python
    py::buffer_info buffer_info = inArray.request();

    // extract data an shape of input array
    T *data = static_cast<T *>(buffer_info.ptr);
    std::vector<ssize_t> shape = buffer_info.shape;
    std::vector<ssize_t> shape_out = {shape[2], shape[3]};

    // 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<T, 4>> A(data, shape[0], shape[1], shape[2], shape[3]);

    // build result tensor with reverse ordering
    Eigen::Tensor<T, 2> out_tensor(shape[2], shape[3]);
    for (int i=0; i < shape_out[0]; i++) {
        for (int j=0; j < shape_out[1]; j++) {
            out_tensor(i, j) = A(0, 0, i, j) * (A(1, 1, i, j) * A(2, 2, i, j) - A(2, 1, i, j) * A(1, 2, i, j))
                             - A(0, 1, i, j) * (A(1, 0, i, j) * A(2, 2, i, j) - A(2, 0, i, j) * A(1, 2, i, j))
                             + A(0, 2, i, j) * (A(1, 0, i, j) * A(2, 1, i, j) - A(2, 0, i, j) * A(1, 1, i, j));
        }

    }

    // return numpy array wrapping eigen tensor's pointer
    return py::array_t<T>(shape_out, 
                          out_tensor.data());  // data pointer
}

PYBIND11_MODULE(SIGNATURE, m) {
    m.def("calculateDeterminant", &calculateDeterminant<double>, py::return_value_policy::move,
        py::arg("inArray"));
}

which somehow gives an incorrect answer for a test array of shape (3,3,1000,4).

Could anyone point me to what could be going wrong?

Of course, when I throw a single 3x3 array it works!!

To be complete, here is how I would achieve this in python:

from numpy import random, zeros_like

def det(A):
    detA = zeros_like(A[0, 0])
    detA = A[0, 0] * (A[1, 1] * A[2, 2] -
                    A[1, 2] * A[2, 1]) -\
        A[0, 1] * (A[1, 0] * A[2, 2] -
                    A[1, 2] * A[2, 0]) +\
        A[0, 2] * (A[1, 0] * A[2, 1] -
                    A[1, 1] * A[2, 0])
    return detA

F = random.random((3,3,1000,4))
detF = det(F)

[YannickJadoul]

@bhaveshshrimali In which way is the answer incorrect? Contrary to @JonasHarsch's example from above, you are not passing strides, and Eigen seems to be column-major/Fortran-style by default. Is that the issue?

[bhaveshshrimali]

@YannickJadoul ,
Thanks for your reply and the link. If I am understanding the strides along each axes correctly then I expect this should work, right? (I tried the other combinations too, namely {shape[0] * sizeof(T), sizeof(T)})

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <unsupported/Eigen/CXX11/Tensor>

namespace py = pybind11;

// return 4 dimensional ndarray with determinant broadcasted across trailing axis
template<class T>
py::array_t<T> calculateDeterminant(py::array_t<T> inArray) {

    // request a buffer descriptor from Python
    py::buffer_info buffer_info = inArray.request();

    // extract data an shape of input array
    T *data = static_cast<T *>(buffer_info.ptr);
    std::vector<ssize_t> shape = buffer_info.shape;
    std::vector<ssize_t> shape_out = {shape[2], shape[3]};

    // 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<T, 4>> A(data, shape[0], shape[1], shape[2], shape[3]);

    // build result tensor with reverse ordering
    Eigen::Tensor<T, 2> out_tensor(shape[2], shape[3]);
    for (int i=0; i < shape_out[0]; i++) {
        for (int j=0; j < shape_out[1]; j++) {
            out_tensor(i, j) = A(0, 0, i, j) * (A(1, 1, i, j) * A(2, 2, i, j) - A(2, 1, i, j) * A(1, 2, i, j))
                             - A(0, 1, i, j) * (A(1, 0, i, j) * A(2, 2, i, j) - A(2, 0, i, j) * A(1, 2, i, j))
                             + A(0, 2, i, j) * (A(1, 0, i, j) * A(2, 1, i, j) - A(2, 0, i, j) * A(1, 1, i, j));
        }

    }

    // return numpy array wrapping eigen tensor's pointer
    return py::array_t<T>(shape_out,
                          {shape[1] * sizeof(T), sizeof(T)},
                          out_tensor.data());  // data pointer
}

PYBIND11_MODULE(eigen, m) {
    m.def("calculateDeterminant", &calculateDeterminant<double>, py::return_value_policy::move,
        py::arg("inArray"));
}

Test

from numpy import zeros_like, random, allclose
from eigen import calculateDeterminant

def det(A):
    detA = zeros_like(A[0, 0])
    detA = A[0, 0] * (A[1, 1] * A[2, 2] -
                    A[1, 2] * A[2, 1]) -\
        A[0, 1] * (A[1, 0] * A[2, 2] -
                    A[1, 2] * A[2, 0]) +\
        A[0, 2] * (A[1, 0] * A[2, 1] -
                    A[1, 1] * A[2, 0])
    return detA

arr = random.random((3,3,1000,4)) 
detCpp = calculateDeterminant(arr)
detPy = det(arr)

assert allclose(detCpp, detPy) # fails

arr2 = random.random((3,3,1,1))
assert allclose( calculateDeterminant(arr2), det(arr2)) # passes

Even though the latter (with arr2) is a degenerate case, I wonder why it works and the former doesn't
Let me know if you notice something off with the above code. Thanks!!

[JonasBreuling]

I messed up with row and column majored storage layout in my previous example. So thank you very much for pointing out this error with this example! Below you find the perfect working code with correct ordering on the eigen side, due to numpy's default storage layout, see here. There where three mistakes in your code

(1) wrong input tensor storage layout
(2) wrong output tensor storage layout
(3) correct strides for the output numpy array / use py::array_t<T, py::array::c_style> without specifying the strides, see comment of @YannickJadoul

eigen.cpp

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <unsupported/Eigen/CXX11/Tensor>

namespace py = pybind11;

// return 4 dimensional ndarray with determinant broadcasted across trailing axis
template<class T>
py::array_t<T> calculateDeterminant(py::array_t<T> inArray) {

    // request a buffer descriptor from Python
    py::buffer_info buffer_info = inArray.request();

    // extract data an shape of input array
    T *data = static_cast<T *>(buffer_info.ptr);
    std::vector<ssize_t> shape = buffer_info.shape;
    std::vector<ssize_t> shape_out = {shape[2], shape[3]};

    // 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<T, 4, Eigen::RowMajor>> A(data, shape[0], shape[1], shape[2], shape[3]); // (1) corrected storage layout of input tensor

    // build result tensor with reverse ordering
    Eigen::Tensor<T, 2, Eigen::RowMajor> out_tensor(shape[2], shape[3]); // (2) corrected storage layout of output tensor
    for (int i=0; i < shape_out[0]; i++) {
        for (int j=0; j < shape_out[1]; j++) {
            out_tensor(i, j) = A(0, 0, i, j) * (A(1, 1, i, j) * A(2, 2, i, j) - A(1, 2, i, j) * A(2, 1, i, j))
                             - A(0, 1, i, j) * (A(1, 0, i, j) * A(2, 2, i, j) - A(1, 2, i, j) * A(2, 0, i, j))
                             + A(0, 2, i, j) * (A(1, 0, i, j) * A(2, 1, i, j) - A(1, 1, i, j) * A(2, 0, i, j));
        }

    }

    // return numpy array wrapping eigen tensor's pointer
    return py::array_t<T>(shape_out,
                          {shape_out[1] * sizeof(T), sizeof(T)}, // (3) corrected strides 
                          out_tensor.data());  // data pointer

    // as YannickJadoul  pointed out an even easier solution is
    /*
    return py::array_t<T, py::array::c_style>(shape_out,
                          out_tensor.data());  // data pointer
    */
}

PYBIND11_MODULE(eigen, m) {
    m.def("calculateDeterminant", &calculateDeterminant<double>, py::return_value_policy::move,
        py::arg("inArray"));
}

test.py

import numpy as np
from eigen import calculateDeterminant

def det(A):
    detA = np.zeros_like(A[0, 0])
    detA = A[0, 0] * (A[1, 1] * A[2, 2] - A[1, 2] * A[2, 1]) - \
        A[0, 1] * (A[1, 0] * A[2, 2] - A[1, 2] * A[2, 0]) + \
        A[0, 2] * (A[1, 0] * A[2, 1] - A[1, 1] * A[2, 0])
    return detA

arr = np.random.random((3,3,10,4)) 
detCpp = calculateDeterminant(arr)
detPy = det(arr)

assert np.allclose(detCpp, detPy) # passes
diff = detCpp - detPy
print(f'detCpp:\n{detCpp}')
print(f'detPy:\n{detPy}')
print(f'diff:\n{diff}')

arr2 = np.random.random((3,4,1,1))
assert np.allclose( calculateDeterminant(arr2), det(arr2)) # passes

[bhaveshshrimali]

There where three mistakes in your code

(1) wrong input tensor storage layout
(2) wrong output tensor storage layout
(3) correct strides for the output numpy array

Thanks a ton @JonasHarsch!! (3) was pretty sloppy on my part :-)

[YannickJadoul]

Thanks, @JonasHarsch!

Minor addition: if you use py::array_t<T, py::array::f_style> (or c_style), you don't need to pass the strides, for a contiguous array. pybind11 will automatically generate the C- or Fortran-style strides.

[JonasBreuling]

Thanks, @YannickJadoul, that's a really convenient way I wasn't aware of! I've adapted the solution with this option.