typeguard warnings, on not having any annotations, now name the corre…

…ct function.

jaxtyping/_decorator.py

@@ -393,16 +393,18 @@ def wrapped_fn(*args, **kwargs):  # pyright: ignore
             param_signature = full_signature.replace(
                 return_annotation=inspect.Signature.empty
             )
+            name = getattr(fn, "__name__", "<no name found>")
+            qualname = getattr(fn, "__qualname__", "<no qualname found>")
             module = getattr(fn, "__module__", "generated")
 
+            # Use the same name so that typeguard warnings look correct.
             full_fn, output_name = _make_fn_with_signature(
-                "check_return", full_signature, module, output=True
+                name, qualname, module, full_signature, output=True
             )
-            full_fn = typechecker(full_fn)
-
             param_fn = _make_fn_with_signature(
-                "check_params", param_signature, module, output=False
+                name, qualname, module, param_signature, output=False
             )
+            full_fn = typechecker(full_fn)
             param_fn = typechecker(param_fn)
 
             @ft.wraps(fn)
@@ -565,17 +567,19 @@ def _check_dataclass_annotations(self, typechecker):
         values[field.name] = value
 
     signature = inspect.Signature(parameters)
-    module = self.__class__.__module__
     f = _make_fn_with_signature(
-        self.__class__.__name__, signature, module, output=False
+        self.__class__.__name__,
+        self.__class__.__qualname__,
+        self.__class__.__module__,
+        signature,
+        output=False,
     )
-    f.__qualname__ = self.__class__.__qualname__
     f = jaxtyped(f, typechecker=typechecker)
     f(self, **values)
 
 
 def _make_fn_with_signature(
-    name: str, signature: inspect.Signature, module: str, output: bool
+    name: str, qualname: str, module: str, signature: inspect.Signature, output: bool
 ):
     """Dynamically creates a function `fn` with name `name` and signature `signature`.
 
@@ -697,6 +701,7 @@ def _make_fn_with_signature(
     exec(fnstr, scope)
     fn = scope[name]
     fn.__module__ = module
+    fn.__qualname__ = qualname
     assert fn is not None
     if output:
         return fn, output_name
@@ -745,7 +750,7 @@ def _get_problem_arg(
         assert keep_annotation is not sentinel
         new_signature = inspect.Signature(new_parameters)
         fn = _make_fn_with_signature(
-            "check_single_arg", new_signature, module, output=False
+            "check_single_arg", "check_single_arg", module, new_signature, output=False
         )
         fn = typechecker(fn)  # but no `jaxtyped`; keep the same environment.
         try:

test/test_array.py

@@ -154,7 +154,7 @@ def g(x: Shaped[Array, "a b"]) -> Shaped[Array, "a b"]:
     g(jnp.array([[1, 2], [3, 4]], dtype=jnp.int8))
     g(jnp.array([[1, 2], [3, 4]], dtype=jnp.uint4))
     g(jnp.array([[1, 2], [3, 4]], dtype=jnp.uint16))
-    g(jr.normal(getkey(), (3, 4), dtype=jnp.complex128))
+    g(jr.normal(getkey(), (3, 4), dtype=jnp.complex64))
     g(jr.normal(getkey(), (3, 4), dtype=jnp.bfloat16))
 
     with pytest.raises(ParamError):

test/test_decorator.py

@@ -96,23 +96,23 @@ def test_context(getkey):
 
 def test_varargs(jaxtyp, typecheck):
     @jaxtyp(typecheck)
-    def f(*args):
+    def f(*args) -> None:
         pass
 
     f(1, 2)
 
 
 def test_varkwargs(jaxtyp, typecheck):
     @jaxtyp(typecheck)
-    def f(**kwargs):
+    def f(**kwargs) -> None:
         pass
 
     f(a=1, b=2)
 
 
 def test_defaults(jaxtyp, typecheck):
     @jaxtyp(typecheck)
-    def f(x: int, y=1):
+    def f(x: int, y=1) -> None:
         pass
 
     f(1)

test/test_generators.py

@@ -20,7 +20,7 @@ def gen(x: Float[Array, "*"]) -> Iterator[Float[Array, "*"]]:
         yield x
 
     @jaxtyp(typecheck)
-    def foo():
+    def foo() -> None:
         next(gen(jnp.zeros(2)))
         next(gen(jnp.zeros((3, 4))))
 
@@ -81,7 +81,7 @@ def g(x: Shaped[torch.Tensor, "*"]) -> Iterator[Shaped[torch.Tensor, "*"]]:
         yield x
 
     @jaxtyp(typecheck)
-    def f():
+    def f() -> None:
         next(g(torch.zeros(1)))
         next(g(torch.zeros(2)))