Expose nb_ndarray th and ndarray_handle owner

[raplonu]

Hey !

I'm working with managed container. My type expose a capsule, a pure C++ type that hold a pointer alongside a reference count.

Creating a nanobind::capsule from this and giving it to a nanobind::ndarray is straightforward. This setup allows the python type to safely utilize the buffer while keeping it alive. However, doing the opposite present some challenges.

First, once created, nanobind doesn't expose ndarray_handle struct preventing me to access the owner member of this type. A potential solution could be to add an accessor function:

PyObject *get_owner(ndarray_handle *th) noexcept;

A more intricate issue arises when trying to retrieve the capsule from a Python type. When calling ndarray_import, owner is always set to nullptr, even if a capsule exists. Nevertheless, the capsule may be found (under base.obj when using numpy, not sure for the other frameworks though). The remaining issue is that to access the capsule, I need to know the internal nb_ndarray type which expose ndarray_handle *th. One solution could be:

ndarray_handle *get_ndrray_handle(handle h);

This function would attempt to cast h to nb_ndarray and then return the th member.

Alternatively, we could perform this in ndarray_import and set owner if found.

Doing that would allow me the check is the capsule is mine and if so, just increment the ref count. Or just copy the ndarray into my own capsule to keep it alive.

What do you think of this ? Thank you !

[wjakob]

It's not obvious how a Python owner of a DLPack tensor could be exposed since only tensors created by nanobind itself are guaranteed to have that property. That information is available when we are returning a tensor. When receiving a tensor in a function binding, we do not know this anymore. All that is left is the DLPack data structure itself, which is intentionally 100% opaque.

[wjakob]

For example, suppose that you wrap MyStruct into a ndarray with a capsule that keeps it alive. The owner field is then that capsule object. Upon returning the tensor to Python, it will likely be converted to an nd-array in some framework such as NumPy.

When you then call another nanobind function binding with this tensor, the owner field will be the NmPy array, not your original capsule wrapping the MyStruct. Of course, these objects still exist, and reference counting ensures that everything is kept alive. But the representation is opaque, and so you can't just walk down a chain of pointers to get to the original object.

[raplonu]

Thanks a lot for your answer! I agree, it seems quite difficult to implement that. I will passe on trying to get back my own capsule back.

However, I still wonder if there is a way to know if there is something to put in my capsule to keep the buffer alive.
I'm quite unfamiliar with DLPack. Does it have any concept of ownership (like numpy OWNDATA) ? Thanks again.