Type stub integral and floating points bindings to accept the corresponding numpy types

[AaronSartin]

If I bind float my_func(float x) I want the type stubs to be my_func(x: float | np.float32 | np.float64) -> float. I can pass an np.float32 at run time without issue, but isinstance(np.float32, float) evaluates false so type checkers throw an error.

I see exactly what I need to change here:

nanobind/include/nanobind/nb_cast.h

Line 187 in 9b3afa9

NB_TYPE_CASTER(T, const_name<std::is_integral_v<T>>("int", "float"))

In my own personal project I'm hacking in the necessary numpy imports with __prefix__ statements, but I think I could figure out how to fix the required imports properly.

@wjakob are you interested in making the typing changes part of the master branch? We can avoid users to have numpy installed with if TYPE_CHECKING blocks around the import statement. If not, do you have any recommendations for how I override the type_caster? I know I can copy all of nb_cast.h and make my own version, but maintaining my own fork of that file doesn't sound like a great solution.

[hpkfft]

Note that the C++ function float my_func(float x) takes and returns single precision, but Python float and NumPy np.float64 are double precision. So, I'm not sure your suggested type stub is ideal....
If you do not specify nb::arg().noconvert(), the double precision Python values will be converted to single precision to be passed to your function from Python, but in general, that is a value-changing conversion that loses precision. Why not also list np.float16? That can be converted to single precision without changing value.

I'm not at all knowledgeable about the Python typing stuff. Maybe it's good that a typing tool complains about automatic conversions that may hurt accuracy? Again, I don't know the right answer.

Though, I would consider preferring the C++ function double my_func(double x)

[AaronSartin]

Note that the C++ function float my_func(float x) takes and returns single precision, but Python float and NumPy np.float64 are double precision. So, I'm not sure your suggested type stub is ideal....

I think that's unrelated to the suggestion. nanobind is already accepting a python float and casting the double down to four bytes.

Why not also list np.float16

I have no objections to adding float16 to the list; it's not a type I commonly use and it wasn't top of mind.

Maybe it's good that a typing tool complains about automatic conversions that may hurt accuracy?

For a simple float interface that's not something nanobind currently deals with and not something I think would be a major value add. If you're talking about a noconvert eigen/numpy interface I think it's a different conversation. The current Eigen typing is only through ArrayLike with useful information on the Annotated dict that isn't actually used by static analyzers.

I think no matter where you go with "should we be stricter on precision conversions?" this example points to "np.float32 should be an accepted input for float my_func(float x)". I think there's an interesting extension here where you can argue that we should return a np.float32 to cut some conversion costs but that requires the user have numpy installed. I would like that for my own projects but I think it's best to leave that conversation for a separate thread given the much larger scope.

[AaronSartin]

Expanding on the desired typing:

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    try:
        # using a small set of types to keep it simple; float16, int16, etc can be include
        # I think nanobind currrently accepts uint*s for int fields as well but I'd have to double check that claim
        from numpy import float32, float64, int32, int64
    except ImportError:
        float32 = float  # type: ignore[assignment, misc]
        float64 = float  # type: ignore[assignment, misc]
        int32 = int  # type: ignore[assignment, misc]
        int64 = int  # type: ignore[assignment, misc]

# to be most correct we should include shape information the `npt.NDArray` stub
# pybind also supported automatic decimal.Decimal -> float conversion; I haven't checked nanobind's behavior yet.
FloatLike = float | float32 | float64
FloatArrayLike = npt.NDArray[float32] | npt.NDArray[float64]

# pybind also supported automatic decimal.Decimal -> int conversion but it was a truncation so I never included it in the type hint.
IntLike = int | int32 | int64
IntArrayLike = npt.NDArray[int32 ] | npt.NDArray[int64]

# probably want a UIntLike too


def func(x: FloatLike) -> float:
    """Automatically coerces float like input but always returns a float."""


def func_array(x: FloatArrayLike) -> npt.NDArray[np.float32]:
    """Automatically coerces a numpy array of float like data to return a float32 array."""


def func_array_no_convert(x: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
    """Automatically coerces a numpy array of float like data to return a float32 array."""

If this sort of behavior is popular, I'd be happy to take a stab at making pr. I'm not familiar enough with the other array interops (e.g. pytorch and tensorflow) to know if they deserve similar treatment.

[hpkfft]

In your own personal project, https://nanobind.readthedocs.io/en/latest/typing.html#pattern-files might be a good solution.

I would suggest that the general problem of defining a type system for number-like types is not something that can or should be undertaken by nanobind.

For those interested in element types of Python arrays, the following may be of interest:
https://data-apis.org/array-api/latest/API_specification/type_promotion.html
https://docs.jax.dev/en/latest/type_promotion.html

[AaronSartin]

I would suggest that the general problem of defining a type system for number-like types is not something that can or should be undertaken by nanobind.

This was my concern; I personally disagree, but it's a defensible position.

In your own personal project, https://nanobind.readthedocs.io/en/latest/typing.html#pattern-files might be a good solution.

Please correct me if I'm misunderstanding but I don't think pattern files work for my use case. I think nanobind only regexes against the qualified name of the function or class and not against the types of the function. I do agree that I could use nb::sig to override signatures but in a large project that would be painful to implement and maintain.

I was interested in overriding the type_caster definition in nb_cast.h but forking that file similarly seems likely to be a headache to maintain. Maybe it would make sense to wrap the type_caster definition in a macro if users want to override it?

[hpkfft]

Sorry, I don't have a good answer for you.
I don't even know what the signatures should be, much less the best way to achieve them.
My thinking is that a type checker is more useful if it warns about value-changing conversions (e.g., narrowing from double to single precision), but I acknowledge that not everyone would like that. But if the type checker accepts anything that Python itself would accept, then what's the point?
I'd like to see a number-like type system that accepts the promotions given in the links above (i.e., allows the forward arrows and warns on everything else).
The numeric tower allows promotions of int to float and it allows promotions of either float or int to complex. That's it. All the NumPy, JAX, PyTorch "dtypes" of array objects deserve a solution, but it's not an easy problem. If I were to tackle it, I would probably take inspiration from the C++ floating-point conversion rank definition and propose something similar to the Python array API standard (link from above) or author a PEP.
As you mentioned above, "ArrayLike with useful information on the Annotated dict that isn't actually used by static analyzers" is not a great solution.

[AaronSartin]

But if the type checker accepts anything that Python itself would accept, then what's the point?

I want the type checker to represent what the binded functions accept. Currently they do accept an np.float64 input but mypy tells me that it's an error. I don't like having to pepper my code with # type: ignore. In the float vs np.float64 vs np.float32 there isn't a meaningful runtime impact; the ffi is much more expensive than the conversion (which is already being done).

I appreciate your desire to have stricter typing on conversions but I think achieving that solely through typing doesn't make a lot of sense. If we wanted that behavior I would suggest that we do something like py::arg("some_float").no_convert() -> my_func(some_float: np.float32) -> .... That is not currently a feature of the library though.

If there is interest in making that a feature I think your points make sense because they would make my suggestion redundant. However, I have a suspicion that disallowing those conversions would be a painful change for many users and not desirable. I personally would like that change but I doubt the users of the binded packages I maintain would agree.

All the NumPy, JAX, PyTorch "dtypes" of array objects deserve a solution

If following through with this suggestion I do think it's important to make sure their dtypes are similarly handled.

[hpkfft]

Those interested in this discussion and in ndarray may also be interested in #442