Pretty error messages: fixes #6.

Phew, this ended up being a pretty complicated change!
The basic summary is that we now support the syntax
```
@jaxtyped(typechecker=typechecker)
def f(...): ...
```
and when using this, we now get pretty error messages about what went
wrong.

(
The old syntax, i.e.
```
@jaxtyped
@typechecker
def f(...): ...
```
is still supported, but doesn't give much information.
)

The internals of this do quite a lot of magic! In particular we
dynamically create quite a lot of functions and test the provided
arguments against their signatures. The overhead should still be
minimal under `jax.jit`, though.
(TODO: what's the overhead like in non-jit situations, e.g. PyTorch?
I've tried to minimise the overhead throughout just to be sure, but
perhaps PyTorch users should stick to the old syntax?)

test/conftest.py

@@ -41,6 +41,29 @@ def typecheck(request):
     return request.param
 
 
+@pytest.fixture(params=(False, True))
+def jaxtyp(request):
+    import jaxtyping
+
+    if request.param:
+        # New-style
+        # @jaxtyping.jaxtyped(typechecker=typechecker)
+        # def f(...)
+        return lambda typechecker: jaxtyping.jaxtyped(typechecker=typechecker)
+    else:
+        # Old-style
+        # @jaxtyping.jaxtyped
+        # @typechecker
+        # def f(...)
+        def impl(typechecker):
+            def decorator(fn):
+                return jaxtyping.jaxtyped(typechecker(fn))
+
+            return decorator
+
+        return impl
+
+
 @pytest.fixture()
 def getkey():
     def _getkey():

test/test_array.py

@@ -33,7 +33,6 @@ from jaxtyping import (
     Bool,
     Float,
     Float32,
-    jaxtyped,
     PRNGKeyArray,
     Shaped,
 )
@@ -41,9 +40,8 @@ from jaxtyping import (
 from .helpers import ParamError, ReturnError
 
 
-def test_basic(typecheck):
-    @jaxtyped
-    @typecheck
+def test_basic(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
     def g(x: Shaped[Array, "..."]):
         pass
 
@@ -82,16 +80,14 @@ def test_dtypes():
             assert key == val.__name__
 
 
-def test_return(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_return(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float[Array, "b c"]) -> Float[Array, "c b"]:
         return jnp.transpose(x)
 
     g(jr.normal(getkey(), (3, 4)))
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def h(x: Float[Array, "b c"]) -> Float[Array, "b c"]:
         return jnp.transpose(x)
 
@@ -99,9 +95,8 @@ def test_return(typecheck, getkey):
         h(jr.normal(getkey(), (3, 4)))
 
 
-def test_two_args(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_two_args(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Shaped[Array, "b c"], y: Shaped[Array, "c d"]):
         return x @ y
 
@@ -109,8 +104,7 @@ def test_two_args(typecheck, getkey):
     with pytest.raises(ParamError):
         g(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (5, 4)))
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def h(x: Shaped[Array, "b c"], y: Shaped[Array, "c d"]) -> Shaped[Array, "b d"]:
         return x @ y
 
@@ -119,9 +113,8 @@ def test_two_args(typecheck, getkey):
         h(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (5, 4)))
 
 
-def test_any_dtype(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_any_dtype(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Shaped[Array, "a b"]) -> Shaped[Array, "a b"]:
         return x
 
@@ -136,14 +129,12 @@ def test_any_dtype(typecheck, getkey):
         g(jr.normal(getkey(), (1,)))
 
 
-def test_nested_jaxtyped(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_nested_jaxtyped(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "b c"], transpose: bool) -> Float32[Array, "c b"]:
         return h(x, transpose)
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def h(x: Float32[Array, "c b"], transpose: bool) -> Float32[Array, "b c"]:
         if transpose:
             return jnp.transpose(x)
@@ -157,9 +148,8 @@ def test_nested_jaxtyped(typecheck, getkey):
         g(jr.normal(getkey(), (2, 3)), False)
 
 
-def test_nested_nojaxtyped(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_nested_nojaxtyped(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "b c"]):
         return h(x)
 
@@ -171,9 +161,8 @@ def test_nested_nojaxtyped(typecheck, getkey):
         g(jr.normal(getkey(), (2, 3)))
 
 
-def test_isinstance(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_isinstance(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "b c"]) -> Float32[Array, " z"]:
         y = jnp.transpose(x)
         assert isinstance(y, Float32[Array, "c b"])
@@ -187,9 +176,8 @@ def test_isinstance(typecheck, getkey):
     g(jr.normal(getkey(), (2, 3)))
 
 
-def test_fixed(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_fixed(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(
         x: Float32[Array, "4 5 foo"], y: Float32[Array, " foo"]
     ) -> Float32[Array, "4 5"]:
@@ -204,9 +192,8 @@ def test_fixed(typecheck, getkey):
         g(c, b)
 
 
-def test_anonymous(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_anonymous(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "foo _"], y: Float32[Array, " _"]):
         pass
 
@@ -215,9 +202,8 @@ def test_anonymous(typecheck, getkey):
     g(a, b)
 
 
-def test_named_variadic(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_named_variadic(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(
         x: Float32[Array, "*batch foo"],
         y: Float32[Array, " *batch"],
@@ -240,8 +226,7 @@ def test_named_variadic(typecheck, getkey):
     with pytest.raises(ParamError):
         g(a2, b1, c)
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def h(x: Float32[Array, " foo *batch"], y: Float32[Array, " foo *batch bar"]):
         pass
 
@@ -255,9 +240,8 @@ def test_named_variadic(typecheck, getkey):
         h(b, c)
 
 
-def test_anonymous_variadic(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_anonymous_variadic(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "... foo"], y: Float32[Array, " foo"]):
         pass
 
@@ -277,9 +261,8 @@ def test_anonymous_variadic(typecheck, getkey):
         g(a3, c)
 
 
-def test_broadcast_fixed(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_broadcast_fixed(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "#4"]):
         pass
 
@@ -290,9 +273,8 @@ def test_broadcast_fixed(typecheck, getkey):
         g(jr.normal(getkey(), (3,)))
 
 
-def test_broadcast_named(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_broadcast_named(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, " #foo"], y: Float32[Array, " #foo"]):
         pass
 
@@ -314,9 +296,8 @@ def test_broadcast_named(typecheck, getkey):
         g(b, a)
 
 
-def test_broadcast_variadic_named(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_broadcast_variadic_named(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, " *#foo"], y: Float32[Array, " *#foo"]):
         pass
 
@@ -373,9 +354,8 @@ def test_broadcast_variadic_named(typecheck, getkey):
         g(o, a)
 
 
-def test_variadic_mixed_broadcast1(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_variadic_mixed_broadcast(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def f(x: Float[Array, " *foo"], y: Float[Array, " #*foo"]):
         pass
 
@@ -389,9 +369,8 @@ def test_variadic_mixed_broadcast1(typecheck, getkey):
     f(c, d)
 
 
-def test_variadic_mixed_broadcast2(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_variadic_mixed_broadcast2(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def f(x: Float[Array, " *#foo"], y: Float[Array, " *foo"]):
         pass
 
@@ -405,9 +384,8 @@ def test_variadic_mixed_broadcast2(typecheck, getkey):
     f(c, d)
 
 
-def test_variadic_mixed_broadcast3(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_variadic_mixed_broadcast3(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def f(
         x: Float[Array, "*B L D"],
         *,
@@ -427,24 +405,20 @@ def test_no_commas():
         Float32[Array, "foo, bar"]
 
 
-def test_symbolic(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_symbolic(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def make_slice(x: Float32[Array, " dim"]) -> Float32[Array, " dim-1"]:
         return x[1:]
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def cat(x: Float32[Array, " dim"]) -> Float32[Array, " 2*dim"]:
         return jnp.concatenate([x, x])
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def bad_make_slice(x: Float32[Array, " dim"]) -> Float32[Array, " dim-1"]:
         return x
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def bad_cat(x: Float32[Array, " dim"]) -> Float32[Array, " 2*dim"]:
         return jnp.concatenate([x, x, x])
 
@@ -464,9 +438,8 @@ def test_symbolic(typecheck, getkey):
         bad_cat(x)
 
 
-def test_incomplete_symbolic(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_incomplete_symbolic(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def foo(x: Float32[Array, " 2*dim"]):
         pass
 
@@ -573,9 +546,8 @@ def test_py310_unions():
     assert isinstance(x, get_args(y))
 
 
-def test_key(typecheck):
-    @jaxtyped
-    @typecheck
+def test_key(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
     def f(x: PRNGKeyArray):
         pass
 
@@ -592,7 +564,7 @@ def test_key(typecheck):
         f(jnp.array(3.0))
 
 
-def test_extension(typecheck, getkey):
+def test_extension(jaxtyp, typecheck, getkey):
     X = Shaped[Array, "a b"]
     Y = Shaped[X, "c d"]
     Z = Shaped[Array, "c d a b"]
@@ -601,8 +573,7 @@ def test_extension(typecheck, getkey):
     X = Float[Array, "a"]
     Y = Float[X, "b"]
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def f(a: X, b: Y):
         ...
 

test/test_decorator.py

@@ -1,9 +1,12 @@
 import abc
 
 import jax.random as jr
+import pytest
 
 from jaxtyping import Array, Float, jaxtyped
 
+from .helpers import ParamError, ReturnError
+
 
 class M(metaclass=abc.ABCMeta):
     @jaxtyped
@@ -85,3 +88,77 @@ def test_context(getkey):
     with jaxtyped("context"):
         assert isinstance(a, Float[Array, "foo bar"])
         assert not isinstance(b, Float[Array, "foo"])
+
+
+def test_varargs(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(*args):
+        pass
+
+    f(1, 2)
+
+
+def test_varkwargs(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(**kwargs):
+        pass
+
+    f(a=1, b=2)
+
+
+def test_defaults(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(x, y=1):
+        pass
+
+    f(1)
+
+
+class _GlobalFoo:
+    pass
+
+
+def test_global_stringified_annotation(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(x: "_GlobalFoo") -> "_GlobalFoo":
+        return x
+
+    f(_GlobalFoo())
+
+    @jaxtyp(typecheck)
+    def g(x: int) -> "_GlobalFoo":
+        return x
+
+    @jaxtyp(typecheck)
+    def h(x: "_GlobalFoo") -> int:
+        return x
+
+    with pytest.raises(ReturnError):
+        g(1)
+
+    with pytest.raises(ParamError):
+        h(1)
+
+
+# This test does not use `jaxtyp(typecheck)` because typeguard does some evil stack
+# frame introspection to try and grab local variables.
+def test_local_stringified_annotation(typecheck):
+    class LocalFoo:
+        pass
+
+    @jaxtyped(typechecker=typecheck)
+    def f(x: "LocalFoo") -> "LocalFoo":
+        return x
+
+    f(LocalFoo())
+
+    @jaxtyped
+    @typecheck
+    def g(x: "LocalFoo") -> "LocalFoo":
+        return x
+
+    g(LocalFoo())
+
+    # We don't check that errors are raised if it goes wrong, since we can't usually
+    # resolve local type annotations at runtime. Best we can hope for is not to raise
+    # a spurious error about not being able to find the type.

test/test_messages.py

@@ -0,0 +1,108 @@
+from typing import Any
+
+import equinox as eqx
+import jax.numpy as jnp
+import pytest
+
+from jaxtyping import Array, Float, jaxtyped, PyTree, TypeCheckError
+
+
+def test_arg_localisation(typecheck):
+    @jaxtyped(typechecker=typecheck)
+    def f(x: str, y: str, z: int):
+        pass
+
+    matches = [
+        "Type-check error whilst checking the parameters of f",
+        "The problem arose whilst typechecking argument 'z'.",
+        r"Called with args: \('hi', 'bye', 'not-an-int'\)",
+        "Called with kwargs: {}",
+        r"Parameter annotations: \(x: str, y: str, z: int\).",
+    ]
+    for match in matches:
+        with pytest.raises(TypeCheckError, match=match):
+            f("hi", "bye", "not-an-int")
+
+    @jaxtyped(typechecker=typecheck)
+    def g(x: Float[Array, "a b"], y: Float[Array, "b c"]):
+        pass
+
+    x = jnp.zeros((2, 3))
+    y = jnp.zeros((4, 3))
+    matches = [
+        "Type-check error whilst checking the parameters of g",
+        "The problem arose whilst typechecking argument 'y'.",
+        r"Called with args: \(f32\[2,3\],\)",
+        r"Called with kwargs: {'y': f32\[4,3\]}",
+        (
+            r"Parameter annotations: \(x: Float\[Array, 'a b'\], y: "
+            r"Float\[Array, 'b c'\]\)."
+        ),
+        "The current values for each jaxtyping axis annotation are as follows.",
+        "a=2",
+        "b=3",
+    ]
+    for match in matches:
+        with pytest.raises(TypeCheckError, match=match):
+            g(x, y=y)
+
+
+def test_return(typecheck):
+    @jaxtyped(typechecker=typecheck)
+    def f(x: PyTree[Any, " T"], y: PyTree[Any, " S"]) -> PyTree[Any, "T S"]:
+        return "foo"
+
+    x = (1, 2)
+    y = {"a": 1}
+    matches = [
+        "Type-check error whilst checking the return value of f",
+        r"Called with args: \(\(1, 2\),\)",
+        r"Called with kwargs: {'y': {'a': 1}}",
+        "Return value: 'foo'",
+        r"Return annotation: PyTree\[Any, \"T S\"\].",
+        (
+            "The current values for each jaxtyping PyTree structure annotation are as "
+            "follows."
+        ),
+        r"T=PyTreeDef\(\(\*, \*\)\)",
+        r"S=PyTreeDef\({'a': \*}\)",
+    ]
+    for match in matches:
+        with pytest.raises(TypeCheckError, match=match):
+            f(x, y=y)
+
+
+def test_dataclass_attribute(typecheck):
+    @jaxtyped(typechecker=typecheck)
+    class M(eqx.Module):
+        x: Float[Array, " *foo"]
+        y: PyTree[Any, " T"]
+        z: int
+
+    x = jnp.zeros((2, 3))
+    y = (1, (3, 4))
+    z = "not-an-int"
+
+    matches = [
+        "Type-check error whilst checking the parameters of M",
+        "The problem arose whilst typechecking argument 'z'.",
+        r"Called with args: \(\)",
+        (
+            r"Called with kwargs: {'x': f32\[2,3\], 'y': \(1, \(3, 4\)\), "
+            r"'z': 'not-an-int'}"
+        ),
+        (
+            r"Parameter annotations: \(x: Float\[Array, '\*foo'\], "
+            r"y: PyTree\[Any, \"T\"\], z: int\)."
+        ),
+        "The current values for each jaxtyping axis annotation are as follows.",
+        r"foo=\(2, 3\)",
+        (
+            "The current values for each jaxtyping PyTree structure annotation are as "
+            "follows."
+        ),
+        r"T=PyTreeDef\(\(\*, \(\*, \*\)\)\)",
+    ]
+    for match in matches:
+        with pytest.raises(TypeCheckError, match=match):
+            M(x, y, z)

test/test_pytree.py

@@ -26,7 +26,7 @@ import jax.random as jr
 import pytest
 
 import jaxtyping
-from jaxtyping import Array, Float, jaxtyped, PyTree
+from jaxtyping import Array, Float, PyTree
 
 from .helpers import make_mlp, ParamError
 
@@ -93,9 +93,8 @@ def test_nested_pytrees(getkey, typecheck):
         g([1, 2, make_mlp()])
 
 
-def test_pytree_array(typecheck):
-    @jaxtyped
-    @typecheck
+def test_pytree_array(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
     def g(x: PyTree[Float[jnp.ndarray, "..."]]):
         pass
 
@@ -107,9 +106,8 @@ def test_pytree_array(typecheck):
         g(1.0)
 
 
-def test_pytree_shaped_array(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_pytree_shaped_array(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: PyTree[Float[jnp.ndarray, "b c"]]):
         pass
 
@@ -196,9 +194,8 @@ def test_subclass_pytree():
     assert not issubclass(int, PyTree)
 
 
-def test_structure_match(typecheck):
-    @jaxtyped
-    @typecheck
+def test_structure_match(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
     def f(x: PyTree[int, " T"], y: PyTree[str, " T"]):
         pass
 
@@ -209,9 +206,8 @@ def test_structure_match(typecheck):
         f(1, ("hi",))
 
 
-def test_structure_prefix(typecheck):
-    @jaxtyped
-    @typecheck
+def test_structure_prefix(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
     def f(x: PyTree[int, " T"], y: PyTree[str, "T ..."]):
         pass
 
@@ -228,9 +224,8 @@ def test_structure_prefix(typecheck):
         f((3, 4, 5), {"a": ("hi", "bye")})
 
 
-def test_structure_suffix(typecheck):
-    @jaxtyped
-    @typecheck
+def test_structure_suffix(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
     def f(x: PyTree[int, " T"], y: PyTree[str, "... T"]):
         pass
 
@@ -245,9 +240,8 @@ def test_structure_suffix(typecheck):
         f((3, 4, 5), {"a": ("hi", "bye")})
 
 
-def test_structure_compose(typecheck):
-    @jaxtyped
-    @typecheck
+def test_structure_compose(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
     def f(x: PyTree[int, " T"], y: PyTree[int, " S"], z: PyTree[str, "S T"]):
         pass
 
@@ -262,8 +256,7 @@ def test_structure_compose(typecheck):
     with pytest.raises(ParamError):
         f((1, 2), {"a": 3}, ({"a": "hi"}, {"a": "bye"}))
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def g(x: PyTree[int, " T"], y: PyTree[int, " S"], z: PyTree[str, "T S"]):
         pass
 
@@ -274,7 +267,7 @@ def test_structure_compose(typecheck):
 
 
 @pytest.mark.parametrize("variadic", (False, True))
-def test_treepath_dependence_function(variadic, typecheck, getkey):
+def test_treepath_dependence_function(variadic, jaxtyp, typecheck, getkey):
     if variadic:
         jtshape = "*?foo"
         shape = (2, 3)
@@ -282,8 +275,7 @@ def test_treepath_dependence_function(variadic, typecheck, getkey):
         jtshape = "?foo"
         shape = (4,)
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def f(
         x: PyTree[Float[Array, jtshape], " T"], y: PyTree[Float[Array, jtshape], " T"]
     ):
@@ -312,7 +304,7 @@ def test_treepath_dependence_dataclass(variadic, typecheck, getkey):
         jtshape = "?foo"
         shape = (4,)
 
-    @jaxtyping._decorator._jaxtyped_typechecker(typecheck)
+    @jaxtyping.jaxtyped(typechecker=typecheck)
     class A(eqx.Module):
         x: PyTree[Float[Array, jtshape], " T"]
         y: PyTree[Float[Array, jtshape], " T"]
@@ -331,9 +323,8 @@ def test_treepath_dependence_dataclass(variadic, typecheck, getkey):
         A((x1, x2), (y2, y1))
 
 
-def test_treepath_dependence_missing_structure_annotation(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_treepath_dependence_missing_structure_annotation(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def f(x: PyTree[Float[Array, "?foo"], " T"], y: PyTree[Float[Array, "?foo"]]):
         pass
 
@@ -343,9 +334,8 @@ def test_treepath_dependence_missing_structure_annotation(typecheck, getkey):
         f(x1, y1)
 
 
-def test_treepath_dependence_multiple_structure_annotation(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_treepath_dependence_multiple_structure_annotation(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def f(x: PyTree[PyTree[Float[Array, "?foo"], " S"], " T"]):
         pass