[QUESTION]: returning `nb::ndarray`

[senyai]

Hello! Just started using nanobind.

• Is there a way to allocate nb::ndarray using nanobind?
• Is it possible to return the same type of nb::ndarray as was passed in, i.e. functionality similar to np.empty_like with exact library (numpy, torch, etc.)? An example in documentation would be nice.
• I'm getting this error with the following code RuntimeError: nanobind::detail::ndarray_export(): copy failed: AttributeError: 'PyCapsule' object has no attribute 'copy'. (linux, python 3.13). Is the problem on my side?

#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>

namespace nb = nanobind;

using Image2D = nb::ndarray<float, nb::shape<-1, -1>, nb::device::cpu>;

Image2D test(Image2D &image) {
  Image2D output(nullptr, {10, 30});
  return output;
}

NB_MODULE(libadrt, m) { m.def("test", &test); }

Thanks

[wjakob]

The last error should be fixed on master, can you try there?

It should be easy to make an empty-like function on your end, but nanobind does not provide one. It needs to deal with many different device types, and providing allocation routines for CUDA, AMD, Intel GPUs, TPUs, etc., would not be practical.

[senyai]

'2.7.1-dev1' - still the same error. I have numpy=1.26.4 if it matters.

I was thinking, I could wrap my "test" call in python,

def test(image):
    match t.__class__.__module__:
        case 'numpy':
            import numpy as np
            buf = np.empty_like(image)
        case 'torch':
            import torch
            buf = torch.empty_like(image)
        case module_name:
            raise NotImplementerError(module_name)
    libadrt_.test(image, buf)
    return buf

Is it an okay solution? (I didn't try it, but is must solve PyCapsule issue)

It should be easy to make an empty-like function on your end, but nanobind does not provide one. It needs to deal with many different device types, and providing allocation routines for CUDA, AMD, Intel GPUs, TPUs, etc., would not be practical.

Could you please update documentation with that information? Preferably in preamble.

[hpkfft]

Arsiniy,

Please see https://nanobind.readthedocs.io/en/latest/ndarray.html#returning-arrays-from-c-to-python
and https://nanobind.readthedocs.io/en/latest/ndarray.html#data-ownership

Having nanobind return DLPack and converting to a specific framework in Python can be a good solution.
See: https://nanobind.readthedocs.io/en/latest/ndarray.html#array-libraries

For example:

#include <cstring>  // memset
#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>

namespace nb = nanobind;

using Image2D = nb::ndarray<float, nb::shape<-1, -1>, nb::device::cpu>;

class MyArray {
    Image2D image;

 public:
    MyArray(std::size_t rows, std::size_t cols) {
        float *data = new float[rows * cols];
        std::memset(data, 0, rows * cols * sizeof(float));
        nb::capsule owner(data, [](void *p) noexcept { delete[] (float *) p; });
        image = Image2D(data, {rows, cols}, std::move(owner));
    }

    Image2D dlpack() { return image; }

    auto dlpack_device() {
        return std::make_pair(nb::device::cpu::value, 0);
    }
};

MyArray test() {
  MyArray output(3, 5);
  return output;
}

NB_MODULE(my_extension, m) {
    nb::class_<MyArray>(m, "MyArray")
        .def("__dlpack__", &MyArray::dlpack)
        .def("__dlpack_device__", &MyArray::dlpack_device);
    m.def("test", &test);
}

Then, in Python,

>>> import numpy as np
>>> import my_extension as me
>>> d = me.test()
>>> d
<my_extension.MyArray object at 0x7f0fda596f60>
>>> a = np.from_dlpack(d)
>>> a
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]], dtype=float32)