Support for ArrayLike objects

[jdumas]

Hi,

I was wondering what's the best way to support ArrayLike as input objects in nanobind bindings? Using a generic nanobind::object() doesn't give hints about the expected type of an argument so it's not ideal.

[wjakob]

There is nanobind::ndarray to check that an input object is an array, and to optionally check that it satisfies various other requirements. This is not just a type annotation--the library checks that the input type indeed satisfies those requirements (otherwise, the called C++ code might crash -- e.g., if it expects an array with a particular layout and is given a different one).

Nanobind also converts this into a type signature, but not a standards compliant one. The state of affairs here is unsatisfactory because the official type toolkits (MyPy extensions, numpy.typing) are really limited and can't absorb all of the metadata that nanobind wishes to provide. If you have ideas on how to improve this, I'm all ears.

[jdumas]

What I would like to do is be able to silently convert an Array-like input of the right size to a NDArray or Eigen::Matrix argument. Consider the following example:

#include <nanobind/eigen/dense.h>
#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>
#include <nanobind/stl/array.h>

#include <Eigen/Dense>

namespace nb = nanobind;

using namespace nb::literals;

using Vec3f = Eigen::Vector<float, 3>;
using NDVec3f = nb::ndarray<float, nb::shape<3>, nb::device::cpu>;

NB_MODULE(nanobind_example_ext, m)
{
    m.def(
        "cross_norm_array",
        [](const std::array<float, 3>& a, const std::array<float, 3>& b) {
            Vec3f a_vec(a[0], a[1], a[2]);
            Vec3f b_vec(b[0], b[1], b[2]);
            return a_vec.cross(b_vec).norm();
        },
        "a"_a,
        "b"_a);

    m.def(
        "cross_norm_eigen",
        [](const Vec3f& a, const Vec3f& b) { return a.cross(b).norm(); },
        "a"_a,
        "b"_a);

    m.def(
        "cross_norm_ndarray",
        [](NDVec3f a, NDVec3f b) {
            Vec3f a_vec(a(0), a(1), a(2));
            Vec3f b_vec(b(0), b(1), b(2));
            return a_vec.cross(b_vec).norm();
        },
        "a"_a,
        "b"_a);
}

And python code:

def test_list():
    a = [1, 2, 3]
    b = [4, 5, 6]
    print(m.cross_norm_array(a, b))
    print(m.cross_norm_eigen(np.array(a), np.array(b)))
    print(m.cross_norm_ndarray(np.array(a), np.array(b)))


def test_tuple():
    a = (1, 2, 3)
    b = (4, 5, 6)
    print(m.cross_norm_array(a, b))
    print(m.cross_norm_eigen(np.array(a), np.array(b)))
    print(m.cross_norm_ndarray(np.array(a), np.array(b)))

Currently I have to call np.array(a) when calling cross_norm_eigen and cross_norm_ndarray, and I'd like to avoid that to simplify the C++ code. The type caster for the cross_norm_array binding already does this silent conversion. Is it possible to achieve with the existing Eigen wrappers in Nanobind?

[wjakob]

This is basically #327, right? This would be absolutely possible--somebody just has to do it, and it isn't on my priority list atm.

Let me explain what this would entail:

• nanobind uses the DLPack and Buffer (only for NumPy) protocols to exchange data with other frameworks. Both are similar, except that DLPack is a bit more general and supports GPU-resident ndarrays and is even usable in packages that don't have a python component at all.

• With this approach, a producer owns some region of memory, and it writes metadata describing the type and layout of the data into a C-style struct which nanobind parses. A user extension may then consume the ndarray, which creates a view and performs some reference counting in the back to ensure that the memory region remains alive until the consumer is done with it.

• The case you are referring to is where the owner of this data isn't a classic nd-array framework like NumPy or PyTorch. There is neither a buffer protocol-compatible object nor a DLPack capsule that nanobind can interpret. Yet, the ArrayLike data could in principle be turned into a packed nd-array.

• The easiest way to solve this problem would be to convert the input Python object into something that can provide the buffer protocol or DLPack metadata on the fly. This can then use all of the remaining infrastructure.

• ..but this should be done without NumPy, since that isn't always present. An alternative would be the builtin array.array (https://docs.python.org/3/library/array.html#module-array) that is, however, limited to depth-1 arrays. So one cannot pass a matrix that way.

• Yet another option would be to go fully custom, parse the input, convert it into packed memory block, and then describe that.

• All of these options involve one major question: what should be the type of the imported array? A Python int could map to all sorts of different signed and unsigned integer types. A Python float could map to half, single, or double precision float. Ideally the implementation could use the numeric type T underlying a value's nd::ndarray<T> type to express a preference to turn the Python ArrayLike into a packed representation with this numeric type.

[jdumas]

Thanks for the detailed explanation. I don't think I'll be able to contribute such a feature at the moment, but this should be helpful if someone has the bandwidth to do it. I think #327 is asking specifically about 1D list, but I think the question is more general (ND list or tuple to NDarray and/or Eigen matrices). Conversion to the target numeric type is perfectly acceptable in this scenario imho.