nanobind uses uses unversioned `DLManagedTensor` which lacks writable flag

[lahwaacz]

Problem description

nanobind uses an "unversioned" managed_dltensor struct which corresponds to the legacy DLManagedTensor:

nanobind/src/nb_ndarray.cpp

Lines 10 to 14 in 879bca4

	struct managed_dltensor {
	dlpack::dltensor dltensor;
	void *manager_ctx;
	void (*deleter)(managed_dltensor *);
	};

This lacks the flags attribute that is present in the DLManagedTensorVersioned struct:

https://github.com/dmlc/dlpack/blob/7f393bbb86a0ddd71fde3e700fc2affa5cdce72d/include/dlpack/dlpack.h#L349-L360

An important flag is DLPACK_FLAG_BITMASK_READ_ONLY which indicates if the array is read-only or writable.

Since numpy/numpy#28600 NumPy's from_dlpack function sets readonly = 1 for unversioned DLPacks.

[wjakob]

This is currently not on my TODO list, and I consider it to be a feature request rather than an bug. Would you be willing to add versioned dlpack tensors to nanobind?

[lahwaacz]

I can look into it, but I'm very new to dlpack...

[wjakob]

Closing as this is a feature request. I am open to pull requests that add this feature.

[hpkfft]

This feature has been added to the development version of nanobind in the git master branch. (#1175)
If you (anybody reading this) are able to test it, feedback would be welcome. (Problems can be reported by adding a comment here if that is more convenient for you.)

[lahwaacz]

@hpkfft Thanks for implementing it! 🎉

I've noticed that when I return a nb::ndarray<> directly from a custom __dlpack__ method, then numpy.from_dlpack(my_array).flags.writeable is False. On the other hand, if I use the new array_api_t and its __dlpack__ method, then the writeable flag is True as expected. It is the trick from

nanobind/tests/test_ndarray.cpp

Lines 544 to 551 in babec16

	.def("__dlpack__", [](nb::pointer_and_handle<MyArray> self,
	nb::kwargs kwargs) {
	using array_api_t = nb::ndarray<nb::array_api, double>;
	nb::object aa = nb::cast(array_api_t(self.p->data(), {5}),
	nb::rv_policy::reference_internal,
	self.h);
	return aa.attr("__dlpack__")(**kwargs);
	})

Actually even if I don't pass the **kwargs then the writeable flag is False. Can you clarify the difference?

[hpkfft]

If you write

      .def("__dlpack__", [](const MyArray& self) {
              return nb::ndarray<double>(self.data(), {5});
           }, nb::rv_policy::reference_internal);

then nanobind is creating a capsule for an unversioned DLManagedTensor, which you are returning.
Neither the code above nor nanobind examines any keyword arguments passed to this __dlpack__() function.
This is how the old DLPack protocol worked, and for backwards compatibility, no changes were made to nanobind here.
The framework annotations are described in https://nanobind.readthedocs.io/en/latest/ndarray.html#returning-arrays-from-c-to-python

Using a recent version of NumPy, numpy.from_dlpack(my_array) calls the __dlpack__() function above with keyword argument max_version set to the tuple (1, 0), which indicates that NumPy can understand the DLManagedTensorVersioned struct. However, an unversioned DLManagedTensor is returned to NumPy. That struct lacks a flag to indicate whether the data is read-only, so NumPy conservatively marks it read-only. (Other frameworks make their own decisions and may assume the data is writable because it probably is.)

When the framework nb::array_api is used, nanobind does not return a capsule. It returns a Python object that has a __dlpack__() function. That function examines keyword arguments. If max_version is (1, 0) or greater, it makes a capsule for a DLManagedTensorVersioned struct. Otherwise (i.e., max_version is less than (1, 0), or is None, or isn't passed at all) then a capsule for an unversioned DLManagedTensor is made.

The implementation of the __dlpack__() function in https://nanobind.readthedocs.io/en/latest/ndarray.html#array-libraries, which you quoted above, forwards all the keyword arguments to the __dlpack__() function of the Python object that is created by nanobind when the nb::array_api is used. That function examines the keyword arguments and honors max_version.

[lahwaacz]

That makes sense, thanks for the explanation!