jaxtyped is missing type annotations

[andriusstank]

Hello,

I'm trying to use this library together with mypy, and ran into a problem that @jaxtyped decorator drops type annotations of the function, effectively disabling static type checking. For example:

from jaxtyping import jaxtyped
from typeguard import typechecked

@jaxtyped(typechecker=typechecked)
def f1(x: int) -> int:
    return x

@typechecked
def f2(x: int) -> int:
    return x

reveal_type(f1)
reveal_type(f2)

$ mypy test2.py 
test2.py:12: note: Revealed type is "Any"
test2.py:13: note: Revealed type is "def (x: builtins.int) -> builtins.int"

@typechecked decorated function retains its type, while @jaxtyped is considered as Any.

I have found a similar issue #217, it was closed, but I didn't really understand why.

[patrick-kidger]

Thanks for the report! Looks like pyright was smart enough to figure things out without the extra annotations (which is why I never noticed), but that mypy needed a helping hand. Should be fixed in #255