[OTHER]: Array Stubgen

[finnBsch]

Hi,

Thanks for the amazing work! In my project, I am returning arrays (e.g. on cuda/pytorch), defined like this

template<typename T>
using GPU3DArray = nb::ndarray<nb::pytorch, T, nb::shape<-1, -1, -1>, nb::c_contig, nb::device::cuda>;

from C++ to python. The generated stubs look like this:

Annotated[ArrayLike, dict(dtype='float32', shape=(None, None, None), order='C', device='cuda', writable=False)]

This does not allow my IDE to understand that it will be a pytorch array and thus it cannot type hint pytorch function for this return value. Is there any way to change this?

According to the documentation, it should already do this, but it doesn't for me. I am using nanobind 2.2.0.

Thanks!