Returning Eigen::Matrix as multi-dimensional ndarray

[eduardohenriquearnold]

Hi, thanks for this awesome project, it's my first time using nanobind and the journey has been very smooth due to the awesome documentation and easy-to-follow API!

I'm trying to bind a class method that returns an Eigen::Matrix. This matrix represents a RGBD image (4 dimensional array with shape H, W, 4). Unfortunately, since Eigen doesn't officially supports multi-dimensional arrays, I'm flattening the array into a single-dimensional 4*H*W array.

Although I can bind this function successfully using the nanobind-provided Eigen type-casters, I would like to return a Numpy array with the correct shape, as opposed to the flattened array. I could do this easily following the excellent documentation in here.

The problem is that I would like my class method to still return a Eigen::Matrix object, so that it's still usable in C++ land. So an alternative binding would be to create a wrapper function that encapsulates the class method and returns the nb::ndarray. My question is: Is it possible to write this wrapper function outside my C++ class and bind it to my Python class?

You may wonder why I would want to do this - I'm trying to create a self-contained library in C++, and adding python bindings as another module. Ideally I would be able to compile the library itself without nanobind dependencies (which would prevent writing the nanobind wrapper inside the C++ class).

[wjakob]

Based on the specificity of the question, I would say it is outside the scope of this project. Options might be to write wrappers/adapters that you can then bind using nanobind (with nb::cast), or further wrapping code in Python that adapts the shape of the returned numpy arrays (probably the easier route).

[eduardohenriquearnold]

Thanks for the response. I decided to implement my own class to represent an image tensor with underlying storage.

class ImageTensor
{
    public:
        ImageTensor(uint height_, uint width_, uint channels_): height(height_), width(width_), channels(channels_), pixels(height_ * width_ * channels_, 0){};
        ImageTensor(const ImageTensor& other):  height(other.height), width(other.width), channels(other.channels), pixels(other.pixels){};
        ImageTensor(){};
        float& operator()(uint i, uint j, uint k)
        {
            return pixels[i * channels * width + j * channels + k];
        };
        ImageTensor& operator=(const ImageTensor& other)
        {
            height = other.height;
            width = other.width;
            channels = other.channels;
            pixels = other.pixels;
            return *this;
        };

        std::vector<float> pixels;
        uint height, width, channels;
};

And wrote a type_caster following the Eigen matrix caster as follows:

    template <>
    struct type_caster<ImageTensor> 
    {
        using NDArray = nb::ndarray<float, nb::shape<nb::any, nb::any, nb::any>, nb::c_contig, nb::device::cpu>;
        using NDArrayCaster = make_caster<NDArray>;

        // NB_TYPE_CASTER(ImageTensor, NDArrayCaster::Name);
        static constexpr auto Name = NDArrayCaster::Name;
        template <typename T_> using Cast = ImageTensor;

        static handle from_cpp(const ImageTensor &t, rv_policy policy, cleanup_list *cleanup) noexcept {
            ImageTensor *image = new ImageTensor(t);
            capsule owner(image, [](void *p) noexcept {
                delete[] (ImageTensor *) p;
            });

            size_t shape[3] = {image->height, image->width, image->channels};
            return NDArrayCaster::from_cpp(NDArray(image->pixels.data(), 3, shape, owner), policy, cleanup);
        };

        bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept = delete;

    };

I think I'm missing something quite obvious, but when I execute my script from Python I get a PyCapsule object instead of a np.ndarray. I'm not sure what's wrong?

[WKarel]

I think that your specialization of nb::ndarray is missing the target library here, i.e. nb::numpy.
See #188 (reply in thread)

[eduardohenriquearnold]

Yes, that was the problem, it works now, thanks! For completeness, I also had to fix the call to the destructor, so the updated version of the type_caster is:

    template <>
    struct type_caster<ImageTensor> 
    {
        using NDArray = nb::ndarray<nb::numpy, float, nb::shape<nb::any, nb::any, nb::any>, nb::c_contig, nb::device::cpu>;
        using NDArrayCaster = make_caster<NDArray>;

        static constexpr auto Name = NDArrayCaster::Name;
        template <typename T_> using Cast = ImageTensor;

        static handle from_cpp(const ImageTensor &t, rv_policy policy, cleanup_list *cleanup) noexcept {
            ImageTensor *image = new ImageTensor(t);
            capsule owner(image, [](void *p) noexcept {
                delete (ImageTensor *) p;
            });

            size_t shape[3] = {image->height, image->width, image->channels};
            return NDArrayCaster::from_cpp(NDArray(image->pixels.data(), 3, shape, owner), rv_policy::reference, cleanup);
        };

        bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept = delete;

    };