nanobind

nanobind copied to clipboard

This PR overcomes two current limitations with tensors in nanobind when using numpy (it does not affect torch or jax).

• returning tensors from C++ always results in read-only arrays. See #46.
• passing read-only arrays from python results in an error. See #42.

• [x] Allow returning writeable arrays from c++ to python.
• [x] Allow passing read-only arrays to from python to c++.
• [x] clean up tests.
• [ ] Add documentation.

@wjakob I've added an nb::writeable flag in tensor creation so far. Let me know if this is on track.

I can confirm that nb::writeable works in my simple test cases.

There is now a readonly flag that allows read-only numpy arrays to be passed through. It forces the translation through the buffer protocol, and has a static-assert guard so that it should only be combined with numpy.

Both issues are solved, but I wonder if there is a way of combining the readonly/writeable flags. I could not see how to do it as they really serve in different places. For clarity, maybe they should be accepts_readonly and returns_writeable.

Based on the table below, we only need buffer protocal for 2 cases: python readonly to C++ readonly, and C++ readwrite to python readwrite.

C++	Python	C++ -> Python	Python -> C++
RO	RO	dlpack	Buffer Protocol
RO	RW	invalid	dlpack
RW	RO	dlpack	invalid
RW	RW	Buffer Protocol	dlpack

It will be great to use the same flag for both cases. Having both writable and readonly may leads to inconsistency settings. It may be helpful to consider writable a property of the nb::tensor obj rather than the C++ or python buffer it binds to. This may help resolve your semantic concern.

Sorry, I don't understand the table -- it looks misleading. Currently nanobind throws an exception if you send it a RO buffer from python. So neither RO -> RO or RO -> RW from python -> C++ work.

Any readonly flag in C++ (at least for numpy arrays) is simply a warning to the programmer, unless we force the array to be const. Perhaps that is an option, but my c++ templating skills are already stretched!

Maybe @wjakob has some suggestions?

Sorry, I don't understand the table -- it looks misleading. Currently nanobind throws an exception if you send it a RO buffer from python. So neither RO -> RO or RO -> RW from python -> C++ work.

Sorry, I should clarify the table is a summary of the data transfer technologies that (I think) can be used for transferring data under different read/write permissions. (and by python, I meant numpy array.) It is not what nanobind does currently.

In any case, I agree with you that we should combine nb::writable and nb::readonly. However, I don't fully understand your comment that "they really serve in different places". In my view, they are simply describing the read/write permission of the buffer associated with nb::tensor object regardless of whether it is from C++ to python or the other way around.

Sorry for taking a long time to respond -- generally this seems like a reasonable set of changes, though it does seem somewhat specific to NumPy. A similar readonly flag is on the horizon for DLPack (https://github.com/data-apis/array-api/issues/191), so whatever this turns into will have to generalize into passing all kinds of read-only tensors made using different frameworks.

A code style comment: in nanobind, opening braces are generally merged with the previous line -- the PR deviates from this convention.

Hi @wjakob. Thanks for the comments. I'll have a look at the DLPack changes, to try and see how they would fit.

I will close this PR, I don't think it is relevant anymore with recent changes to nanobind. I am happy to reopen it if you disagree.