Support runtime type-checking of generic functions

[davnn]

Hi,

do you think it is conceivable to implement type checks for generics, e.g. generic array types or generic data types, see below.

from jaxtyping import jaxtyped
from beartype import beartype
from jax import Array as JaxArray
from torch import Tensor as TorchArray
from numpy import ndarray as NumpyArray

GenericArray = TypeVar("GenericArray", JaxArray, NumpyArray, TorchArray)
GenericFloat = TypeVar("GenericFloat", Float16, Float32, Float64)

@jaxtyped
@beartype
def f(a: Shaped[GenericArray, "n"]) -> Shaped[GenericArray, "n"]:
    return a

@jaxtyped
@beartype
def f(a: GenericFloat[NumpyArray, "n"]) -> GenericFloat[NumpyArray, "n"]:
    return a

I would be happy to contribute, but I am unsure if there is even a possiblity of success.

[patrick-kidger]

Yup, I think this should be possible! I'd be happy to take a PR on this.