How to make a type_caster for std::mdspan to convert from/to ML framework-agnostic ndarray

[Char-Aznable]

Hi,

I have implemented a custom type_caster for std::mdspan (using rapids/raft's implementation). I wrote this type_caster basing on nanobind's Eigen::Map's type_caster. However, one limitation is that the corresponding ndarray annotation type is hardcoded to the numpy framework. I want to make this std::mdspan type_caster framework agnostic. Can I simply remove the numpy annotation from this line? What alternative approaches can I use?

[wjakob]

It boils down to the question of: what kind of Python object would you like to have? If you remove the numpy annotation, you'll get a DLPack capsule which is not exactly user-friendly.

[Char-Aznable]

I want to make the binding functions taking a mdspan to work with various Python tensors/arrays. E.g.,

using MDSpan = std::mdspan<int, std::extents<int, 2>>;

void myFunc1(MDSpan input) { ...; }
MDSpan myFunc2(MDSpan input) { ...; return MDSpan(...); }

NB_MODULE(my_ext, m) {
    m.def("myFunc1", &myFunc1);
    m.def("myFunc2", &myFunc2);
}
 

Is there a way to make myFunc1 to work with torch, numpy and jax input?
How about myFunct2, which return a tensors/arrays? -- I guess this is only possible if I provide a wrapper/helper class of the return type so that the intended framework is explicit

[wjakob]

You'll have to write your own type caster for this, which you can base on the Eigen one. However, you'll be on your own if you pursue this path -- I unfortunately can't provide this level of support here.

[Char-Aznable]

Sure. I'm no asking any one to implement this into nanobind but I'm asking how I should write the type caster for this. Is making a type_caster for each individual framework annotation by replacing the numpy annotation and duplicating the rest of the code the recommended way to pursue this? I just need some high level ideas to rule out a few possibilities. E.g., am I correct that myFunc1 example above can work with torch, numpy and jax input by mimicking Eigen::Map example? But myFunc2 won't work in any case