Possible for nanobind to support Awkward arrays?

[matthewfeickert]

This is not a discussion about asking if Awkward has any plans to switch to using nanobind over pybind11 — that's a separate discussion that I don't want to force on the team now. This is a question that is based out of my own ignorance of how things work on if the Awkward team thinks there should be any reason that nanobind couldn't support Awkward in the same way that is supports other n-dimensional arrays through the nanobind::ndarray API.

With the release of nanobind v1.0.x, the "The nb::ndarray<..> class" docs mention

nanobind can exchange n-dimensional arrays (henceforth “ndarrays”) with popular array programming frameworks including NumPy, PyTorch, TensorFlow, and JAX. It supports zero-copy exchange using two protocols:

The classic buffer protocol.

• The classic buffer protocol.
• DLPack, a GPU-compatible generalization of the buffer protocol.

To use this feature, you must add the include directive

#include <nanobind/ndarray.h>

to your code. Following this, you can bind functions with nb::ndarray<...>-typed parameters and return values.

which seems to be in contrast with the pybind11 docs on integration with NumPy arrays that only discuss NumPy.

While this is really a question that needs to be put to the nanobind team, does the Awkward team see any technical reason why the list of n-dimensional arrays that nanobind can exchange with in a zero-copy way couldn't also in the future include Awkward?

[agoose77]

@matthewfeickert would you be able to elaborate upon what you'd like to be able to do here?

With Awkward 2, nearly all of Awkward Array is implemented in Python. We have some performance logic (C++ kernels, ArrayBuilder, LayoutBuilder) written in fast compiled code, but most of what makes Awkward "tick" is interpreted.

As such, I'm not sure what a nanobind interface would look like, or offer to the user. For regular arrays, it would be possible to convert their layouts to e.g. nb::ndarray, but this would potentially fail at runtime if the layout is not regular. I don't know if that's a huge amount more convenient than either requiring a NumPy array as the argument (e.g. my_func(ak.to_numpy(x))) or invoking Awkward from C++ (https://nanobind.readthedocs.io/en/latest/api_core.html#python-object-api).

We could add support for DLPack (which is a todo item) that would succeed only for regular arrays. That's an agenda item that I've not spoken much to @jpivarski about yet :)

[matthewfeickert]

With Awkward 2, nearly all of Awkward Array is implemented in Python. We have some performance logic (C++ kernels, ArrayBuilder, LayoutBuilder) written in fast compiled code, but most of what makes Awkward "tick" is interpreted.

Yeah, I think this is the critical point of why my original idea fundamentally doesn't make sense. That as well as @jpivarski's follow up point

I agree with @agoose77 in not seeing any benefit to having a special interface to nanobind's nb.ndarray class because Awkward Arrays are not "ndarrays". To get them in that form, one would have to call (explicitly or implicitly) ak.to_numpy or ak.to_cupy, and then nanobind already knows what to do with a NumPy or CuPy array.

So this answers the question. 👍

[jpivarski]

I've been thinking, and have mentioned a few times, that Awkward could possibly replace its pybind11 bindings with nanobind bindings, only because we don't use very much of pybind11's features and can perhaps simplify things, get smaller binaries, and maybe have faster start-up times and call overhead. A few features still use the C++ interface, such as ArrayBuilder, JSON I/O, and AwkwardForth. The kernel functions, which are a much larger fraction of Awkward's functionality, are called through ctypes, not pybind11, and it's unclear to me whether there would be an advantage to generating and compiling nanobind interfaces for all of them, too.

But that's an entirely internal change, and not high on any priority list. It wouldn't be seen from the outside.

I agree with @agoose77 in not seeing any benefit to having a special interface to nanobind's nb.ndarray class because Awkward Arrays are not "ndarrays". To get them in that form, one would have to call (explicitly or implicitly) ak.to_numpy or ak.to_cupy, and then nanobind already knows what to do with a NumPy or CuPy array.

(By the way, we used to use DLPack to deal with CUDA array buffers in Awkward v1—as an internal implementation—but we don't have to anymore because we let CuPy deal with the interface between Python and C++, in analogy with NumPy and C. Presumably, CuPy uses DLPack.)

Basically, we now live one layer higher in abstraction. We don't have C or C++ interfaces, but the libraries we use do. And if you want to get nested structure/jagged arrays between languages, that's what Arrow is for, and that's why we have a lightweight ak.to_arrow function.

[matthewfeickert]

But that's an entirely internal change, and not high on any priority list. It wouldn't be seen from the outside.

Yeah, agreed. That's part of what I meant by

that's a separate discussion that I don't want to force on the team now.
.

(By the way, we used to use DLPack to deal with CUDA array buffers in Awkward v1—as an internal implementation—but we don't have to anymore because we let CuPy deal with the interface between Python and C++, in analogy with NumPy and C. Presumably, CuPy uses DLPack.)

🚀 That's slick and great. :)

Basically, we now live one layer higher in abstraction. We don't have C or C++ interfaces, but the libraries we use do. And if you want to get nested structure/jagged arrays between languages, that's what Arrow is for, and that's why we have a lightweight ak.to_arrow function.

This is a really nice summary in just two sentences. It along with @agoose77's point sums up what I was forgetting last night in terms of how the scope of the awkward-cpp kernels. 👍