NDArray element type constraints -- types are too strict?

[nurpax]

I upgraded to the latest Nanobind @ 10cc510 and my implot stubs stubs started flagging type errors that didn’t appear before. This may not really be a regression in Nanobind but rather that VSCode Pylance was less strict about ArrayLikes than NDArrays. I'm looking for guidance on what's a good way to express NDArray inputs without introducing lot of boilerplate or unnecessary function overloads.

In my implot bindings, I'm declaring plot_line as below:

typedef nb::ndarray<const double, nb::ndim<1>, nb::device::cpu, nb::c_contig> ndarray_1d;

m.def("plot_line", [](const char* label_id, ndarray_1d& xs, ndarray_1d& ys, ImPlotLineFlags_ flags) {
    ImPlot::PlotLine(label_id, (const double*)xs.data(), (const double*)ys.data(), xs.shape(0), flags);
}, "label_id"_a, "xs"_a, "ys"_a, "flags"_a.sig("LineFlags.NONE") = ImPlotLineFlags_None, line_docstring2);

With the older Nanobind version, this produced:

def plot_line(label_id: str, 
    xs: Annotated[ArrayLike, dict(dtype='float64', shape=(None), order='C', device='cpu', writable=False)], 
    ys: Annotated[ArrayLike, dict(dtype='float64', shape=(None), order='C', device='cpu', writable=False)], 
    flags: LineFlags = LineFlags.NONE) -> None:

with the new version, it outputs:

def plot_line(label_id: str, 
    xs: Annotated[NDArray[float64], dict(shape=(None,), order='C', device='cpu', writable=False)], 
    ys: Annotated[NDArray[float64], dict(shape=(None,), order='C', device='cpu', writable=False)], 
    flags: LineFlags = LineFlags.NONE) -> None:

Unfortunately VSCode now complains about the types if I pass a non-float64 array to plot_line(), e.g.

data_x = np.arange(0, 10, 0.25)  # NDArray[floating[Any])
data_y = np.sin(data_x)
implot.plot_line("Sine Wave", data_x, data_y)

(passing a non-double NDArray to this function should be completely fine, as due to the 'double' constraint, the array is converted to doubles by nanobind.)

This outputs a warning:

Argument of type "_Array1D[floating[Any]]" cannot be assigned to parameter "xs" of type "NDArray[float64]" in function "plot_line"
  "ndarray[tuple[int], dtype[floating[Any]]]" is not assignable to "ndarray[_Shape, dtype[float64]]"
    Type parameter "_DType_co@ndarray" is covariant, but "dtype[floating[Any]]" is not a subtype of "dtype[float64]"
      "dtype[floating[Any]]" is not assignable to "dtype[float64]"
        Type parameter "_SCT_co@dtype" is covariant, but "floating[Any]" is not a subtype of "float64"
          "floating[Any]" is not assignable to "float64"P

The xs and ys args are declared as typedef nb::ndarray<const double, nb::ndim<1>, nb::device::cpu, nb::c_contig> which means the input arguments gets automatically converted to a double valued array. This is exactly what I want and it works, it's just that the stub for plot_line() doesn't reflect that. I'd like the xs and ys args to be typed less strict as just NDArray instead of NDArray[np.float64] because that captures the type of inputs accepted by this function.

Below is my current attempt at declaring this in another way that produces a less strict type signature:

typedef nb::ndarray<nb::ndim<1>, nb::device::cpu, nb::c_contig> ndarray_1d_any; // no double constraint!

m.def("plot_line", [](const char* label_id, ndarray_1d_any& xs, ndarray_1d_any& ys, ImPlotLineFlags_ flags) {
    auto xv = xs.view<const double, nb::ndim<1>, nb::device::cpu, nb::c_contig>();
    auto yv = ys.view<const double, nb::ndim<1>, nb::device::cpu, nb::c_contig>();
    ImPlot::PlotLine(label_id, (const double*)xv.data(), (const double*)yv.data(), xv.shape(0), flags);
}, "label_id"_a, "xs"_a, "ys"_a, "flags"_a.sig("LineFlags.NONE") = ImPlotLineFlags_None);

This produces the declaration I'm after (below), but the code doesn't actually work as the .view() calls do not convert the data.

def plot_line(label_id: str, values: Annotated[NDArray, dict(shape=(None,), order='C', device='cpu')], xscale: float = 1.0, xstart: float = 0.0, flags: LineFlags = LineFlags.NONE) -> None:

Any thoughts on what's a good way to express these more relaxed inputs so that the types also reflect what inputs are accepted? IMO the current generated stubs simply are too strict.

[hpkfft]

First, I must mention that I have not generated stubs nor used any python typing tools.
But I want to do so.
Anyway, I think you want your function parameters to be typed as ArrayLike to allow for conversions, and your function return types to be specific, e.g., NDArray[float64].
For function parameters, I guess it makes a difference whether noconvert() is used.

[wjakob]

The expressiveness of type annotations for nd-arrays is frustrating — a lot is lost in translation. I've made the change in commit 10cc510 because this was requested by many users. See #442 and #1094, #442, #983, and #984 for some starting points. This topic has been raised many times.

To move to even more nuanced type annotations, I think we need a concrete proposal and a draft modification to stubgen that implements this proposal. I am not working with NumPy types in any of my projects, so it would be helpful to have a user of such annotations take the initiative.

[nurpax]

I filed #1155. It's a little open ended but I thought it'd be a good place to invite other NDArray commenters to discuss.

[nurpax]

I don't feel like I know NDArray type annotations well enough to write a detailed proposal. In general, I'm fine with switching to NDArray from ArrayLike and my uses would probably be fixed by switching mostly to NDArray[Any]. I wonder if anything would come out of spawning a new GH issue from this and inviting everyone from bugs #442, #1094, #983 and #984 to hash this out?

I was going to suggest that handling .noconvert() args differently from the default and render the stub as NDArray[Any] when there's no .noconvert(). But it's not that simple, as there can be multiple overloads with different constraints and NB will try to match overloads first, and only if there's no match, it will perform conversion. The signatures kind of make sense here.

TBH, I'd be completely happy to "solve" this by manually writing out the type for such args. It's just that this is a little tedious currently, as there's no way to write the type annotation of a single arg, you can only do it for the whole function struct sig for module_::def(), and for functions with complex types and many arguments, writing out the whole thing is cumbersome. I know there's sig() for args, but alas, that can only be just to specify the default value, not the type for that argument.

More concretely, I wouldn't mind doing this (non-working example):

m.def("plot_line", [](const char* label_id, ndarray_1d& xs, ndarray_1d& ys, ImPlotLineFlags_ flags) {
    ImPlot::PlotLine(label_id, (const double*)xs.data(), (const double*)ys.data(), xs.shape(0), flags);
}, "label_id"_a, "xs"_a.sig("NDArray[Any]"), "ys"_a.sig("NDArray[Any]"), "flags"_a.sig("LineFlags.NONE") = ImPlotLineFlags_None);

But right now, to override argument type hints, I'd have to do write it out completely:

m.def("plot_line", [](const char* label_id, ndarray_1d& xs, ndarray_1d& ys, ImPlotLineFlags_ flags) {
    ImPlot::PlotLine(label_id, (const double*)xs.data(), (const double*)ys.data(), xs.shape(0), flags);
}, "label_id"_a, "xs"_a, "ys"_a, "flags"_a = ImPlotLineFlags_None, 
nb::sig("def plot_line(label_id: str, xs: Annotated[NDArray[Any], dict(shape=(None,), order='C', device='cpu', writable=False)], ys: Annotated[NDArray[Any], dict(shape=(None,), order='C', device='cpu', writable=False)], flags: LineFlags = LineFlags.NONE) -> None");