nb::ndarray as a binding instead of type caster?

[zingdle]

now nb::ndarray is a type caster (listed in here), which leads to some unexpected behaviour:

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

namespace nb = nanobind;

NB_MODULE(myext, m) {
  m.def("foo", [](nb::ndarray<> a) {
    auto numpy = nb::module_::import_("numpy");

    // add one dimension
    std::vector<int> shape(a.shape_ptr(), a.shape_ptr() + a.ndim());
    shape.insert(shape.begin(), 1);

    auto b = numpy.attr("broadcast_to")(a, shape);
    nb::print(b);
  });
}

>>> import numpy as np
>>> import myext
>>> a = np.arange(4)
>>> a
array([0, 1, 2, 3])
>>> myext.foo(a)
[[<capsule object "dltensor" at 0x7f674c492f90>
  <capsule object "dltensor" at 0x7f674c492f90>
  <capsule object "dltensor" at 0x7f674c492f90>
  <capsule object "dltensor" at 0x7f674c492f90>]]

in pybind11, equivalent py::array wraps numpy array and provides simpler way to manipulate it on python side directly (listed here), the above code works as expected:

#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

namespace py = pybind11;

PYBIND11_MODULE(myext, m) {
    m.def("foo", [](py::array a) {
        auto numpy = py::module::import("numpy");

        // add one dimension
        std::vector<int> shape(a.shape(), a.shape() + a.ndim());
        shape.insert(shape.begin(), 1);

        auto b = numpy.attr("broadcast_to")(a, shape);
        py::print(b);
    });
}

>>> import numpy as np
>>> import myext
>>> a = np.arange(4)
>>> a
array([0, 1, 2, 3])
>>> myext.foo(a)
[[0 1 2 3]]

I'm curious about the design decision on making nb::ndarray a type caster. As a caster, the shape_ptr() and stride_ptr() point to a newly allocated memory, not the underlying array properties. This means that for a singleton array, everytime I cast it to a nb::ndarray, the shape or stride ptr varies.

Also, would it be more convenient to stick to the underlying element instead of the capsule object "dltensor" when using the nb::ndarray?

[wjakob]

Actually, it's a type caster in both frameworks. However, the design of nb::ndarray is very different because it simultaneously interfaces with different tensor frameworks (including ones that store data on the GPU). If you just have an nb::ndarray and return it, nanobind forgot what framework it came from originally and returns the most general possible representation (a DLPack capsule that can in principle be opened by any framework). Try annotating the type with a framework label, e.g. nb::ndarray<nb::numpy>

[wjakob]

FYI, I also pushed d089d89 that fixes the issue in another way: when returning/casting a nb::ndarray back to python, it remembers the Python object that was originally associated with this array.

[zingdle]

Thanks, I see why nb::ndarray is different.

However, for CPU tensors, it may be convenient to retrived the underlying arrays shape_ptr(), instead of a newly allocated shape vector and control nb::ndarray's lifetime in a shared_ptr way. For example, consider a function that return the CPU tensor's shape ptr. In such case, the created nb::ndarray shape vector get destroyed on stack and the memory is gone.