Add support for PaddlePaddle Tensors for nb::ndarray

[Enigmatisms]

PR Type

Enhancement

Related topics

nb::ndarray

Intro

Hi! I’m a researcher from the PaddlePaddle team. After our conversation at SIGGRAPH Asia 2024 (where you recommended migrating from pybind11 to nanobind), I gave nanobind a try and was impressed by its performance. As part of this effort, I’ve extended nb::ndarray to support paddle.Tensor, allowing direct return of PaddlePaddle tensors from C++ via nb::paddle.

This PR includes:

• ​​New nb::paddle​​: Mirroring nb::pytorch’s behavior, with modifications in nb_ndarray.cpp for consistent tensor handling.
• ​​Unit tests​​: Adapted from existing PyTorch tests (using @needs_torch) to validate PaddlePaddle integration.
• ​​Documentation​​: Added usage guidelines for nb::paddle (preview welcome—I’m unsure how it renders on the website).
• Removed some unnecessary semicolons in the test_ndarray.py.

This change maintains backward compatibility while enabling PaddlePaddle users to leverage nanobind’s efficiency. Looking forward to feedback!

[Enigmatisms]

.cpu() can actually be removed. Other parts are fine.

[Enigmatisms]

Found a problem though, don't know whether it is common for PyTorch and Numpy (I think it is):

For these DL/Math frameworks, the tensors/ndarrays are managed by the framework, therefore, simple conversion might fail. For example (taking Paddle as an example):

paddle::Tensor some_tensor(...);
auto shapes = some_tensor.dims();
auto strides = some_tensor.strides();
# Only viable when the paddle-supporting PR is merged
return nb::ndarray<nb::paddle>(some_tensor.data(), ndim, shapes.data(),
                                           {}, strides.data(), dtype,
                                           device_type, device_id);

The above code will fail, since some_tensor is just a local variable, when the function returns, some_tensor.data() will be invalid since the original owner is GC-ed. The python end will just print random numbers even for zero tensors. deleter of the ndarray must be used, since the owner should be kept alive (so we better pass it out from the function's scope and hold on to it somewhere else). deleter nb::capsule can only have void (*)(void*) noexcept as deleter func, and even if we use lambda, nothing can be captured (since direct capturing will result in compilation failure). The current solution is:

struct LifeTimeManager {
    // steal from Tensor, and help us manage the life time
    explicit LifeTimeManager(paddle::Tensor&& obj) : raii_obj(std::move(obj)) {}
    
    void operator()(void*) const {
        // A magic way: we don't use this, 
    }

    paddle::Tensor raii_obj;
};

struct CapsuleDeleter {
    std::function<void(void*)> func;
    
    CapsuleDeleter(LifeTimeManager&& ltm) 
        : func([ltm=std::move(ltm)](void* p) mutable { ltm(p); }) {}
    
    static void cleanup(void* ptr) noexcept {
        auto* self = static_cast<CapsuleDeleter*>(ptr);
        self->func(self);
        delete self;
    }
};

// upon returning:
auto* deleter_obj = new CapsuleDeleter(LifeTimeManager(std::move(tensor)));
    
nb::capsule deleter(
        deleter_obj, 
        "tensor_deleter", 
        &CapsuleDeleter::cleanup
);

return nb::ndarray<nb::paddle>(some_tensor.data(), ndim, shapes.data(),
                                           deleter, strides.data(), dtype,
                                           device_type, device_id);

Maybe I am overcomplicating things here? Is there any simple way for this? I think this is not a framework-specific problem.