diff --git a/docs/api/array.md b/docs/api/array.md
index fda2dbd..ddbf109 100644
--- a/docs/api/array.md
+++ b/docs/api/array.md
@@ -24,7 +24,7 @@ In addition some modifiers can be applied:
     `def add(x: Float[Array, "#foo"], y: Float[Array, "#foo"]) -> Float[Array, "#foo"]`.
 - Prepend `_` to a dimension to disable any runtime checking of that dimension (so that it can be used just as documentation). This can also be used as just `_` on its own: e.g. `"b c _ _"`.
 - Documentation-only names (i.e. they're ignored by jaxtyping) can be handled by prepending a name followed by `=` e.g. `Float[Array, "rows=4 cols=3"]`.
-- Prepend `?` to a dimension to indicate that its size can vary within a PyTree structure. (See [PyTree annotations](../pytree.md).)
+- Prepend `?` to a dimension to indicate that its size can vary within a PyTree structure. (See [PyTree annotations](../pytree/).)
 
 When using multiple modifiers, their order does not matter.
 
diff --git a/docs/api/runtime-type-checking.md b/docs/api/runtime-type-checking.md
index 772ec29..62adc6e 100644
--- a/docs/api/runtime-type-checking.md
+++ b/docs/api/runtime-type-checking.md
@@ -2,27 +2,40 @@
 
 (See the [FAQ](../faq.md) for details on static type checking.)
 
-Runtime type checking **synergises beautifully with `jax.jit`!** All shape checks will be performed at trace-time only, and will not impact runtime performance.
+Runtime type checking **synergises beautifully with `jax.jit`!** All shape checks will be performed only whilst tracing, and will not impact runtime performance.
 
-Runtime type-checking should be performed using a library like [typeguard](https://github.com/agronholm/typeguard) or [beartype](https://github.com/beartype/beartype).
+There are two approaches: either use [`jaxtyping.jaxtyped`][] to typecheck a single function, or [`jaxtyping.install_import_hook`][] to typecheck a whole codebase.
 
-The types provided by `jaxtyping`, e.g. `Float[Array, "batch channels"]`, are all compatible with `isinstance` checks, e.g. `isinstance(x, Float[Array, "batch channels"])`. This means that jaxtyping should be compatible with all runtime type checkers out-of-the-box.
+In either case, the actual business of checking types is performed with the help of a runtime type-checking library. The two most popular are [beartype](https://github.com/beartype/beartype) and [typeguard](https://github.com/agronholm/typeguard). (If using typeguard, then specifically the version `2.*` series should be used. Later versions -- `3` and `4` -- have some known issues.)
 
-Some additional context is needed to ensure consistency between multiple argments (i.e. that shapes match up between arrays). For this, you can use either `jaxtyping.jaxtyped` to add this capability to a single function, or `jaxtyping.install_import_hook` to add this capability to a whole codebase. If either are too much magic for you, you can safely use neither and have just single-argument type checking.
+---
 
 ::: jaxtyping.jaxtyped
 
 ---
 
-It can be a lot of effort to add `@jaxtyped` decorators all over your codebase.
-(Not to mention that double-decorators everywhere are a bit ugly.)
-
-The easier option is usually to use the import hook.
-
 ::: jaxtyping.install_import_hook
 
 ---
 
+#### Pytest hook
+
+The import hook can be installed at test-time only, as a pytest hook. From the command line the syntax is:
+```
+pytest --jaxtyping-packages=foo,bar.baz,beartype.beartype
+```
+or in `pyproject.toml`:
+```toml
+[tool.pytest.ini_options]
+addopts = "--jaxtyping-packages=foo,bar.baz,beartype.beartype"
+```
+or in `pytest.ini`:
+```ini
+[pytest]
+addopts = --jaxtyping-packages=foo,bar.baz,beartype.beartype
+```
+This example will apply the import hook to all modules whose names start with either `foo` or `bar.baz`. The typechecker used in this example is `beartype.beartype`.
+
 #### IPython extension
 
 If you are running in an IPython environment (for example a Jupyter or Colab notebook), then the jaxtyping hook can be automatically ran via a custom magic:
@@ -32,3 +45,7 @@ import jaxtyping
 %jaxtyping.typechecker beartype.beartype  # or any other runtime type checker
 ```
 Place this at the start of your notebook -- everything that is directly defined in the notebook, after this magic is run, will be hook'd.
+
+#### Other runtime type-checking libraries
+
+Beartype and typeguard happen to be the two most popular runtime type-checking libraries (at least at time of writing), but jaxtyping should be compatible with all runtime type checkers out-of-the-box. The runtime type-checking library just needs to provide a type-checking decorator (analgous to `beartype.beartype` or `typeguard.typechecked`), and perform `isinstance` checks against jaxtyping's types.
diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
index 992e1dd..c0e7396 100644
--- a/jaxtyping/__init__.py
+++ b/jaxtyping/__init__.py
@@ -29,7 +29,7 @@ from ._array_types import (
     has_jax,
     set_array_name_format as set_array_name_format,
 )
-from ._decorator import jaxtyped as jaxtyped
+from ._decorator import jaxtyped as jaxtyped, TypeCheckError as TypeCheckError
 from ._import_hook import install_import_hook as install_import_hook
 from ._ipython_extension import load_ipython_extension as load_ipython_extension
 
diff --git a/jaxtyping/_array_types.py b/jaxtyping/_array_types.py
index 22e48f1..1f118c8 100644
--- a/jaxtyping/_array_types.py
+++ b/jaxtyping/_array_types.py
@@ -27,6 +27,7 @@ from typing import Any, Literal, NoReturn, Optional, Union
 
 import numpy as np
 
+from ._raise import jaxtyping_raise, jaxtyping_raise_from
 from ._storage import get_shape_memo, get_treepath_memo, set_shape_memo
 
 
@@ -124,13 +125,16 @@ def _check_dims(
                 # Make a copy to avoid `__builtins__` getting added as a key.
                 eval_size = eval(cls_dim.expr, single_memo.copy())
             except NameError as e:
-                raise NameError(
-                    f"Cannot process symbolic dimension '{cls_dim.elem_string}' as "
-                    "some dimension names have not been processed. In practice you "
-                    "should usually only use symbolic dimensions in annotations for "
-                    "return types, referring only to dimensions annotated for "
-                    "arguments."
-                ) from e
+                jaxtyping_raise_from(
+                    NameError(
+                        f"Cannot process symbolic dimension '{cls_dim.elem_string}' as "
+                        "some dimension names have not been processed. In practice you "
+                        "should usually only use symbolic dimensions in annotations "
+                        "for return types, referring only to dimensions annotated for "
+                        "arguments."
+                    ),
+                    e,
+                )
             if eval_size != obj_size:
                 return False
         else:
@@ -186,8 +190,8 @@ class _MetaAbstractArray(type):
             if len(repr_dtype) == 2 and repr_dtype[0] == "torch":
                 dtype = repr_dtype[1]
             else:
-                raise RuntimeError(
-                    "Unrecognised array/tensor type to extract dtype from"
+                jaxtyping_raise(
+                    RuntimeError("Unrecognised array/tensor type to extract dtype from")
                 )
 
         if cls.dtypes is not _any_dtype:
@@ -553,10 +557,12 @@ def _make_array(array_type, dim_str, dtypes, name):
 
 class _MetaAbstractDtype(type):
     def __instancecheck__(cls, obj: Any) -> NoReturn:
-        raise RuntimeError(
-            f"Do not use `isinstance(x, jaxtyping.{cls.__name__})`. If you want to "
-            "check just the dtype of an array, then use "
-            f'`jaxtyping.{cls.__name__}[jnp.ndarray, "..."]`.'
+        jaxtyping_raise(
+            RuntimeError(
+                f"Do not use `isinstance(x, jaxtyping.{cls.__name__})`. If you want to "
+                "check just the dtype of an array, then use "
+                f'`jaxtyping.{cls.__name__}[jnp.ndarray, "..."]`.'
+            )
         )
 
     def __getitem__(cls, item: tuple[Any, str]):
@@ -567,6 +573,7 @@ class _MetaAbstractDtype(type):
                 "Ellipsis can be used to accept any shape: `Float[Array, '...']`."
             )
         array_type, dim_str = item
+        dim_str = dim_str.strip()
         del item
         if typing.get_origin(array_type) in _union_types:
             out = [
diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py
index f857c1a..8314cac 100644
--- a/jaxtyping/_decorator.py
+++ b/jaxtyping/_decorator.py
@@ -20,10 +20,11 @@
 import dataclasses
 import functools as ft
 import inspect
+import itertools as it
 import sys
 import types
 import weakref
-from typing import get_args, get_origin
+from typing import Any, get_args, get_origin, get_type_hints, overload
 
 
 try:
@@ -34,113 +35,340 @@ else:
     traceback_util.register_exclusion(__file__)
 
 
-from ._storage import get_shape_memo, pop_shape_memo, push_shape_memo
+from ._storage import pop_shape_memo, push_shape_memo
 
 
 _jaxtyped_fns = weakref.WeakSet()
 
 
-def jaxtyped(fn):
-    """Used in conjunction with a runtime type checker. Decorate a function with this to
-    have shapes checked for consistency across multiple arguments.
+class TypeCheckError(TypeError):
+    pass
 
-    Note that `@jaxtyped` is applied above the type checker.
+
+TypeCheckError.__module__ = "jaxtyping"  # appears in error messages
+
+
+@overload
+def jaxtyped(*, typechecker=None):
+    ...
+
+
+@overload
+def jaxtyped(fn, *, typechecker=None):
+    ...
+
+
+def jaxtyped(fn=None, *, typechecker=None):
+    """Decorate a function with this to perform runtime type-checking of its arguments
+    and return value. Decorate a dataclass to perform type-checking of its attributes.
 
     !!! Example
 
         ```python
         # Import both the annotation and the `jaxtyped` decorator from `jaxtyping`
-        from jaxtyping import Array, Float32, jaxtyped
+        from jaxtyping import Array, Float, jaxtyped
 
         # Use your favourite typechecker: usually one of the two lines below.
         from typeguard import typechecked as typechecker
         from beartype import beartype as typechecker
 
-        # Write your function. @jaxtyped must be applied above @typechecker!
-        @jaxtyped
-        @typechecker
-        def batch_outer_product(x: Float32[Array, "b c1"],
-                                y: Float32[Array, "b c2"]
-                              ) -> Float32[Array, "b c1 c2"]:
+        # Type-check a function
+        @jaxtyped(typechecker=typechecker)
+        def batch_outer_product(x: Float[Array, "b c1"],
+                                y: Float[Array, "b c2"]
+                              ) -> Float[Array, "b c1 c2"]:
             return x[:, :, None] * y[:, None, :]
+
+        # Type-check a dataclass
+        @jaxtyped(typechecker=typechecker)
+        @dataclass
+        class MyDataclass:
+            x: int
+            y: Float[Array "b c"]
         ```
 
+    **Arguments:**
+
+    - `fn`: The function or dataclass to decorate.
+    - `typechecker`: The runtime type-checker to use. This should be a function
+        decorator that will raise an exception if there is a type error, e.g.
+        ```python
+        @typechecker
+        def f(x: int):
+            pass
+
+        f("a string is not an integer")  # this line should raise an exception
+        ```
+        Common choices are `typechecker=beartype.beartype` or
+        `typechecker=typeguard.typechecked`. Can also be set as `typechecker=None` to
+        skip automatic runtime type-checking, but still support manual `isinstance`
+        checks inside the function body:
+        ```python
+        @jaxtyped
+        def f(x):
+            assert isinstance(x, Float[Array, "batch channel"])
+        ```
+
+    **Returns:**
+
+    If `fn` is a function (including a `staticmethod`, `classmethod`, or `property`),
+    then a wrapped function is returned.
+
+    If `fn` is a dataclass, then `fn` is returned directly, and additionally its
+    `__init__` method is wrapped and modified in-place.
+
     **Notes for advanced users**
 
-    Put precisely, all `isinstance` shape checks are scoped to the thread-local dynamic
-    context of a `jaxtyped` call. A new dynamic context will allow different dimensions
-    sizes to be bound to the same name. After this new dynamic context is finished
-    then the old one is returned to.
+    Put precisely, the axis names in e.g. `Float[Array, "batch channels"]` and the
+    structure names in e.g. `PyTree[int, "T"]` are all scoped to the thread-local
+    dynamic context of a `jaxtyped`-wrapped function. A new dynamic context will allow
+    different values to be bound to the same name. After this new dynamic context is
+    finished then the old one is returned to.
 
-    For example, this means you could leave off the `@jaxtyped` decorator to enforce
-    that this function use the same axis sizes as the function it was called from.
+    Binding of a value against a name is done with an `isinstance` check, for example
+    `isinstance(jnp.zeros((3, 4)), Float[Array, "dim1 dim2"])` will bind `dim1=3` and
+    `dim2=4`.
 
-    Likewise, this means you can use `isinstance` checks inside a function body
-    and have them contribute to the same collection of consistency checks performed
-    by a typechecker against its arguments. (Or even forgo a typechecker that analyses
-    arguments, and instead just do your own manual `isinstance` checks.)
+    This means you can use `isinstance` checks inside a function body and have them
+    contribute to the same collection of consistency checks performed by a typechecker
+    against its arguments. (Or even forgo a typechecker that analyses arguments, and
+    instead just do your own manual `isinstance` checks.)
 
-    Only `isinstance` checks that pass will contribute to the store of axis name-size
-    pairs; those that fail will not. As such it is safe to write e.g.
+    Only `isinstance` checks that pass will contribute to the store of values; those
+    that fail will not. As such it is safe to write e.g.
     `assert not isinstance(x, Float32[Array, "foo"])`.
+
+    !!! Info "Old syntax"
+
+        jaxtyping previously (before v0.2.24) recommended using this double-decorator
+        syntax:
+        ```python
+        @jaxtyped
+        @typechecker
+        def f(...): ...
+        ```
+        This is still supported, but the `jaxtyped(typechecker=typechecker)` syntax
+        discussed above will produce easier-to-debug error messages. Under the hood, the
+        new syntax more carefully manipulates the typechecker so as to determine where
+        a type-check error arises.
     """
-    if type(fn) is types.FunctionType and fn in _jaxtyped_fns:
+
+    if fn is None:
+        return ft.partial(jaxtyped, typechecker=typechecker)
+    elif type(fn) is types.FunctionType and fn in _jaxtyped_fns:
+        return fn
+    elif inspect.isclass(fn):
+        if dataclasses.is_dataclass(fn) and typechecker is not None:
+            # This does not check that the arguments passed to `__init__` match the
+            # type annotations. There may be a custom user `__init__`, or a
+            # dataclass-generated `__init__` used alongside
+            # `equinox.field(converter=...)`
+
+            init = fn.__init__
+
+            @ft.wraps(init)
+            def __init__(self, *args, **kwargs):
+                init(self, *args, **kwargs)
+                # `fn.__init__` is late-binding to the `__init__` function that
+                # we're in now. (Or to someone else's monkey-patch.) Either way,
+                # this checks that we're in the "top-level" `__init__`, and not one
+                # that is being called via `super()`. We don't want to trigger too
+                # early, before all fields have been assigned.
+                #
+                # We're not checking `if self.__class__ is fn` because Equinox
+                # replaces the with a defrozen version of itself during `__init__`,
+                # so the check wouldn't trigger.
+                #
+                # We're not doing this check by adding it to the end of the
+                # metaclass `__call__`, because Python doesn't allow you
+                # monkey-patch metaclasses.
+                if self.__class__.__init__ is fn.__init__:
+                    _check_dataclass_annotations(self, typechecker)
+
+            fn.__init__ = __init__
         return fn
-    elif inspect.isclass(fn):  # allow decorators on class definitions
-        if dataclasses.is_dataclass(fn):
-            # TODO(kidger): unify this branch with `_jaxtyped_typechecker` below,
-            # perhaps if we ever do typechecking on an argument-by-argument basis.
-            init = jaxtyped(fn.__init__)
-            fn.__init__ = init
-            return fn
-        else:
-            raise ValueError(
-                "jaxtyped may only be added as a class decorator to dataclasses"
-            )
     # It'd be lovely if we could handle arbitrary descriptors, and not just the builtin
     # ones. Unfortunately that means returning a class instance with a __get__ method,
     # and that turns out to break loads of other things. See beartype issue #211 and
     # jaxtyping issue #71.
     elif isinstance(fn, classmethod):
-        return classmethod(jaxtyped(fn.__func__))
+        return classmethod(jaxtyped(fn.__func__, typechecker=typechecker))
     elif isinstance(fn, staticmethod):
-        return staticmethod(jaxtyped(fn.__func__))
+        return staticmethod(jaxtyped(fn.__func__, typechecker=typechecker))
     elif isinstance(fn, property):
         if fn.fget is None:
             fget = None
         else:
-            fget = jaxtyped(fn.fget)
+            fget = jaxtyped(fn.fget, typechecker=typechecker)
         if fn.fset is None:
             fset = None
         else:
-            fset = jaxtyped(fn.fset)
+            fset = jaxtyped(fn.fset, typechecker=typechecker)
         if fn.fdel is None:
             fdel = None
         else:
-            fdel = jaxtyped(fn.fdel)
+            fdel = jaxtyped(fn.fdel, typechecker=typechecker)
         return property(fget=fget, fset=fset, fdel=fdel)
     elif fn == "context":
         return _JaxtypingContext()
     else:
+        if typechecker is None:
+            # Probably being used in the old style as
+            # ```
+            # @jaxtyped
+            # @typechecker
+            # def foo(x: int): ...
+            # ```
+            # in which case make a best-effort attempt to add shape information for any
+            # type errors.
+            @ft.wraps(fn)
+            def wrapped_fn(*args, **kwargs):  # pyright: ignore
+                memos = push_shape_memo()
+                try:
+                    return fn(*args, **kwargs)
+                except Exception as e:
+                    if sys.version_info >= (3, 11) and _no_jaxtyping_note(e):
+                        shape_info = _exc_shape_info(memos)
+                        if shape_info != "":
+                            msg = (
+                                "The preceding error occurred within the scope of a "
+                                "`jaxtyping.jaxtyped` function, and may be due to a "
+                                "typecheck error. "
+                            )
+                            e.add_note(_jaxtyping_note_str(_spacer + msg + shape_info))
+                    raise
+                finally:
+                    pop_shape_memo()
 
-        @ft.wraps(fn)
-        def wrapped_fn(*args, **kwargs):
-            memos = push_shape_memo()
+        else:
+            # New-style
+            # ```
+            # @jaxtyped(typechecker=typechecker)
+            # def foo(x: int): ...
+            # ```
+            # in which case we can do a better job reporting errors.
+
+            full_signature = inspect.signature(fn)
             try:
-                return fn(*args, **kwargs)
-            except Exception as e:
-                if sys.version_info >= (3, 11) and _no_jaxtyping_note(e):
-                    shape_info = _exc_shape_info(memos)
-                    if shape_info != "":
-                        msg = (
-                            "The preceding error occurred within the scope of a "
-                            "`jaxtyping.jaxtyped` function, and may be due to a "
-                            "typecheck error. "
-                        )
-                        e.add_note(_jaxtyping_note_str(_spacer + msg + shape_info))
-                raise
-            finally:
-                pop_shape_memo()
+                destring_annotations = get_type_hints(fn, include_extras=True)
+            except NameError:
+                # Best-effort attempt to destringify annotations.
+                pass
+            else:
+                new_params = []
+                for p_name, p_value in full_signature.parameters.items():
+                    p_annotation = destring_annotations.get(p_name, p_value.annotation)
+                    p_value = p_value.replace(annotation=p_annotation)
+                    new_params.append(p_value)
+                return_annotation = destring_annotations.get(
+                    "return", full_signature.return_annotation
+                )
+                full_signature = full_signature.replace(
+                    parameters=new_params, return_annotation=return_annotation
+                )
+
+            param_signature = full_signature.replace(
+                return_annotation=inspect.Signature.empty
+            )
+            module = getattr(fn, "__module__", "generated")
+
+            full_fn, output_name = _make_fn_with_signature(
+                "check_return", full_signature, module, output=True
+            )
+            full_fn = typechecker(full_fn)
+
+            param_fn = _make_fn_with_signature(
+                "check_params", param_signature, module, output=False
+            )
+            param_fn = typechecker(param_fn)
+
+            @ft.wraps(fn)
+            def wrapped_fn(*args, **kwargs):
+                # Raise bind-time errors before we do any shape analysis. (I.e. skip
+                # the pointless jaxtyping information for a non-typechecking failure.)
+                param_signature.bind(*args, **kwargs)
+
+                memos = push_shape_memo()
+                try:
+                    # First type-check just the parameters before the function is
+                    # called.
+                    try:
+                        param_fn(*args, **kwargs)
+                    except Exception as e:
+                        if hasattr(e, "_jaxtyping_malformed"):
+                            raise
+                        else:
+                            argmsg = _get_problem_arg(
+                                param_signature, args, kwargs, module, typechecker
+                            )
+                            try:
+                                name = fn.__name__
+                            except AttributeError:
+                                name = fn.__class__.__name__
+                            paramstr = _remove_typing(param_signature)
+                            msg = (
+                                "Type-check error whilst checking the parameters of "
+                                f"{name}.{argmsg}\n"
+                                f"Called with args: {_pformat(args)}\n"
+                                f"Called with kwargs: {_pformat(kwargs)}\n"
+                                f"Parameter annotations: {paramstr}.\n"
+                                + _exc_shape_info(memos)
+                            )
+                            raise TypeCheckError(msg) from e
+
+                    # Actually call the function.
+                    out = fn(*args, **kwargs)
+
+                    if full_signature.return_annotation is not inspect.Signature.empty:
+                        # Now type-check the return value. We need to include the
+                        # parameters in the type-checking here in case there are any
+                        # type variables shared across the parameters and return.
+                        #
+                        # Incidentally this does mean that if `fn` mutates its arguments
+                        # so that they no longer satisfy their type annotations, this
+                        # will throw an error here. But that's like, super weird, so
+                        # don't do that. An error in that scenario is probably still
+                        # desirable.
+                        #
+                        # There is a small performance concern here when used in
+                        # non-jit'd contexts, like PyTorch, due to the duplicate
+                        # checking of the parameters. Unfortunately there doesn't seem
+                        # to be a way around that, so c'est la vie.
+                        kwargs[output_name] = out
+                        try:
+                            full_fn(*args, **kwargs)
+                        except Exception as e:
+                            if hasattr(e, "_jaxtyping_malformed"):
+                                raise
+                            else:
+                                try:
+                                    name = fn.__name__
+                                except AttributeError:
+                                    name = fn.__class__.__name__
+                                paramstr = _remove_typing(param_signature)
+                                returnstr = _remove_typing(
+                                    full_signature.return_annotation
+                                )
+                                if returnstr.startswith(
+                                    "<class '"
+                                ) and returnstr.endswith("'>"):
+                                    returnstr = returnstr[8:-2]
+                                kwargs.pop(output_name)
+                                msg = (
+                                    "Type-check error whilst checking the return value "
+                                    f"of {name}.\n"
+                                    f"Called with args: {_pformat(args)}\n"
+                                    f"Called with kwargs: {_pformat(kwargs)}\n"
+                                    f"Return value: {_pformat(out)}\n"
+                                    f"Parameter annotations: {paramstr}.\n"
+                                    f"Return annotation: {returnstr}.\n"
+                                    + _exc_shape_info(memos)
+                                )
+                                raise TypeCheckError(msg) from e
+
+                    return out
+                finally:
+                    pop_shape_memo()
 
         _jaxtyped_fns.add(wrapped_fn)
         return wrapped_fn
@@ -154,24 +382,19 @@ class _JaxtypingContext:
         pop_shape_memo()
 
 
-@jaxtyped
 def _check_dataclass_annotations(self, typechecker):
+    """Creates and calls a function that checks the attributes of `self`
+
+    `self` should be a dataclass instancae. `typechecker` should be e.g.
+    `beartype.beartype` or `typeguard.typechecked`.
+    """
+    parameters = []
+    values = {}
     for field in dataclasses.fields(self):
-        for kls in self.__class__.__mro__:
-            try:
-                annotation = kls.__annotations__[field.name]
-            except KeyError:
-                pass
-            else:
-                break
-        else:
-            raise TypeError
+        annotation = field.type
         if isinstance(annotation, str):
-            # Don't support stringified annotations. These are basically impossible to
+            # Don't check stringified annotations. These are basically impossible to
             # resolve correctly, so just skip them.
-            # This does mean that annotations like `type["Foo"]` will just fail. There
-            # doesn't seem to be any way to even detect a partially-stringified
-            # annotation.
             continue
         if get_origin(annotation) is type:
             args = get_args(annotation)
@@ -186,100 +409,236 @@ def _check_dataclass_annotations(self, typechecker):
         except AttributeError:
             continue  # allow uninitialised fields, which are allowed on dataclasses
 
-        @typechecker
-        def typecheck(x: annotation):
-            pass
+        parameters.append(
+            inspect.Parameter(
+                field.name,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                annotation=field.type,
+            )
+        )
+        values[field.name] = value
 
-        try:
-            typecheck(value)
-        except Exception as e:
-            if sys.version_info >= (3, 11) and _no_jaxtyping_note(e):
-                shape_info = _exc_shape_info(get_shape_memo())
-                if shape_info != "":
-                    msg = (
-                        "The above typechecking error occurred due to a mismatch "
-                        f"between the value and annotation for field '{field.name}' in "
-                        "dataclass "
-                        f"'{self.__class__.__module__}.{self.__class__.__qualname__}'. "
-                    )
-                    e.add_note(_jaxtyping_note_str(_spacer + msg + shape_info))
-            raise
+    signature = inspect.Signature(parameters)
+    module = self.__class__.__module__
+    f = _make_fn_with_signature(
+        self.__class__.__name__, signature, module, output=False
+    )
+    f = jaxtyped(f, typechecker=typechecker)
+    f(**values)
 
 
-def _jaxtyped_typechecker(typechecker):
-    """A decorator added by the import hook to all classes. Only affects dataclasses.
+def _make_fn_with_signature(
+    name: str, signature: inspect.Signature, module: str, output: bool
+):
+    """Dynamically creates a function `fn` with name `name` and signature `signature`.
 
-    Will be called as
-    ```
-    @_jaxtyped_typechecker(beartype.beartype)
-    @dataclasses.dataclass
-    class SomeDataclass:
-        ...
-    ```
+    If `output=True` then `fn` will consume an additional keyword-only argument (in
+    addition to the provided signature), and will directly return this argument. In this
+    case the returned value from `_make_fn_with_signature` is a 2-tuple of `(fn, name)`,
+    where `fn` is the generated function, and `name` is the name of this extra argument.
 
-    After initialisation, this will check that all fields of the dataclass match their
-    specified type annotation.
+    If `output=False` then `fn` will just have a single `pass` statement, and the
+    returned value from `_make_fn_with_signature` will just be `fn`.
+
+    ---
+
+    Note that this function operates by dynamically creating and eval'ing a string, not
+    simply by assigning `__signature__` and `__annotations__`. The latter is enough for
+    typeguard (at least v2), but does not work with beartype (at least v16).
     """
-    # typechecker is expected to probably be either `typeguard.typechecked`, or
-    # `beartype.beartype`, or `None`.
+    pos = []
+    pos_or_key = []
+    varpos = []
+    key = []
+    varkey = []
+    for p in signature.parameters.values():
+        if p.kind == inspect.Parameter.POSITIONAL_ONLY:
+            pos.append(p)
+        elif p.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD:
+            pos_or_key.append(p)
+        elif p.kind == inspect.Parameter.VAR_POSITIONAL:
+            varpos.append(p)
+        elif p.kind == inspect.Parameter.KEYWORD_ONLY:
+            key.append(p)
+        elif p.kind == inspect.Parameter.VAR_KEYWORD:
+            varkey.append(p)
+        else:
+            assert False
 
-    if typechecker is None:
-        typechecker = lambda x: x
-
-    def _wrapper(kls):
-        assert inspect.isclass(kls)
-        if dataclasses.is_dataclass(kls):
-            # This does not check that the arguments passed to `__init__` match the
-            # type annotations. There may be a custom user `__init__`, or a
-            # dataclass-generated `__init__` used alongside
-            # `equinox.field(converter=...)`
-
-            init = kls.__init__
-
-            @ft.wraps(init)
-            def __init__(self, *args, **kwargs):
-                init(self, *args, **kwargs)
-                # `kls.__init__` is late-binding to the `__init__` function that we're
-                # in now. (Or to someone else's monkey-patch.) Either way, this checks
-                # that we're in the "top-level" `__init__`, and not one that is being
-                # called via `super()`. We don't want to trigger too early, before all
-                # fields have been assigned.
-                #
-                # We're not checking `if self.__class__ is kls` because Equinox replaces
-                # the with a defrozen version of itself during `__init__`, so the check
-                # wouldn't trigger.
-                #
-                # We're not doing this check by adding it to the end of the metaclass
-                # `__call__`, because Python doesn't allow you monkey-patch metaclasses.
-                if self.__class__.__init__ is kls.__init__:
-                    _check_dataclass_annotations(self, typechecker)
-
-            kls.__init__ = __init__
-        return kls
-
-    return _wrapper
-
-
-def _no_jaxtyping_note(e):
-    try:
-        notes = e.__notes__
-    except AttributeError:
-        return True
+    param_names = frozenset(signature.parameters.keys())
+    if output:
+        output_name = _gensym(param_names, prefix="ret")
+        outstr = "return " + output_name
+        param_names = param_names | frozenset({output_name})
+        key.append(inspect.Parameter(output_name, kind=inspect.Parameter.KEYWORD_ONLY))
     else:
-        for note in notes:
-            if isinstance(note, _jaxtyping_note_str):
-                return False
-        return True
+        outstr = "pass"
+
+    scope = {name: None}
+    name_to_annotation = {}
+    name_to_default = {}
+    param_triples = (
+        (p.name, p.annotation, p.default) for p in signature.parameters.values()
+    )
+    if output:
+        triples = it.chain(
+            param_triples,
+            [
+                ("return", signature.return_annotation, inspect.Signature.empty),
+                (output_name, Any, inspect.Signature.empty),
+            ],
+        )
+    else:
+        triples = it.chain(
+            param_triples,
+            [("return", signature.return_annotation, inspect.Signature.empty)],
+        )
+    for p_name, p_annotation, p_default in triples:
+        annotation_name = _gensym(frozenset(scope.keys()) | param_names, prefix="T")
+        name_to_annotation[p_name] = annotation_name
+        if p_annotation is inspect.Signature.empty or isinstance(p_annotation, str):
+            # If we have a stringified annotation here it's because the get_type_hints
+            # lookup above failed. Typically this occurs when using a local variable as
+            # the annotation. In this case we really have no idea what the annotation
+            # refers to, so just set it to Any.
+            # This does mean that we don't handle partially-stringified local
+            # annotations, e.g. `type["Foo"]` for some local type `Foo`. Those will
+            # probably just error out. Nothing better we can do about that
+            # unfortunately.
+            scope[annotation_name] = Any
+        else:
+            scope[annotation_name] = p_annotation
+        default_name = _gensym(frozenset(scope.keys()) | param_names, prefix="default")
+        name_to_default[p_name] = default_name
+        scope[default_name] = p_default
+
+    argstr_pieces = []
+    if len(pos) > 0:
+        for p in pos:
+            argstr_pieces.append(_make_argpiece(p, name_to_annotation, name_to_default))
+        argstr_pieces.append("/")
+    if len(pos_or_key) > 0:
+        for p in pos_or_key:
+            argstr_pieces.append(_make_argpiece(p, name_to_annotation, name_to_default))
+    if len(varpos) == 1:
+        [p] = varpos
+        argstr_pieces.append(
+            "*" + _make_argpiece(p, name_to_annotation, name_to_default)
+        )
+    else:
+        assert len(varpos) == 0
+        if len(key) > 0:
+            argstr_pieces.append("*")
+    if len(key) > 0:
+        for p in key:
+            argstr_pieces.append(_make_argpiece(p, name_to_annotation, name_to_default))
+    if len(varkey) == 1:
+        [p] = varkey
+        argstr_pieces.append(
+            "**" + _make_argpiece(p, name_to_annotation, name_to_default)
+        )
+    else:
+        assert len(varkey) == 0
+    argstr = ", ".join(argstr_pieces)
+
+    if signature.return_annotation is inspect.Signature.empty:
+        retstr = ""
+    else:
+        retstr = f"-> {name_to_annotation['return']}"
+
+    fnstr = f"def {name}({argstr}){retstr}:\n    {outstr}"
+    exec(fnstr, scope)
+    fn = scope[name]
+    fn.__module__ = module
+    assert fn is not None
+    if output:
+        return fn, output_name
+    else:
+        return fn
 
 
-class _jaxtyping_note_str(str):
-    pass
+def _gensym(names: frozenset[str], prefix: str) -> str:
+    assert prefix.isidentifier()
+    output_index = 0
+    output_name = prefix + str(output_index)
+    while output_name in names:
+        output_index += 1
+        output_name = prefix + str(output_index)
+    assert output_name.isidentifier()
+    return output_name
 
 
-_spacer = "--------------------\n"
+def _make_argpiece(p, name_to_annotation, name_to_default):
+    if p.default is inspect.Signature.empty:
+        return f"{p.name}: {name_to_annotation[p.name]}"
+    else:
+        return f"{p.name}: {name_to_annotation[p.name]} = {name_to_default[p.name]}"
+
+
+def _get_problem_arg(
+    param_signature: inspect.Signature, args, kwargs, module, typechecker
+) -> str:
+    """Determines which argument was likely to be the problematic one responsible for
+    raising a type-check error.
+    """
+    # No performance concerns, as this is only used when we're about to raise an error
+    # anyway.
+    for keep_name in param_signature.parameters.keys():
+        new_parameters = []
+        for p_name, p in param_signature.parameters.items():
+            if p_name == keep_name:
+                new_parameters.append(
+                    inspect.Parameter(p.name, p.kind, annotation=p.annotation)
+                )
+            else:
+                new_parameters.append(inspect.Parameter(p.name, p.kind))
+        new_signature = inspect.Signature(new_parameters)
+        fn = _make_fn_with_signature(
+            "check_single_arg", new_signature, module, output=False
+        )
+        fn = typechecker(fn)  # but no `jaxtyped`; keep the same environment.
+        try:
+            fn(*args, **kwargs)
+        except Exception:
+            return f"\nThe problem arose whilst typechecking argument '{keep_name}'."
+    else:
+        # Could not localise the problem to a single argument -- probably due to
+        # e.g. a mismatched typevar, which each individual argument is okay with.
+        return ""
+
+
+def _remove_typing(x):
+    x = str(x)
+    x = x.replace(" jaxtyping.", " ")
+    x = x.replace("[jaxtyping.", "[")
+    x = x.replace("'jaxtyping.", "'")
+    x = x.replace(" typing.", " ")
+    x = x.replace("[typing.", "[")
+    x = x.replace("'typing.", "'")
+    return x
+
+
+def _pformat(x):
+    # No performance concerns from delayed imports -- this is only used when we're about
+    # to raise an error anyway.
+    try:
+        # TODO(kidger): this is pretty ugly. We have a circular dependency
+        # equinox->jaxtyping->equinox. We could consider moving all the pretty-printing
+        # code from equinox into jaxtyping maybe? Or into some shared dependency?
+        import equinox as eqx
+
+        pformat = eqx.tree_pformat
+    except Exception:
+        import pprint
+
+        pformat = ft.partial(pprint.pformat, indent=2, compact=True)
+    return pformat(x)
 
 
 def _exc_shape_info(memos) -> str:
+    """Gives debug information on the current state of jaxtyping's internal memos.
+    Used in type-checking error messages.
+    """
     single_memo, variadic_memo, pytree_memo = memos
     pieces = []
     if len(single_memo) > 0 or len(variadic_memo) > 0:
@@ -294,9 +653,29 @@ def _exc_shape_info(memos) -> str:
                 pieces.append(f"{name}={shape}")
     if len(pytree_memo) > 0:
         pieces.append(
-            "The current values for each jaxtyping pytree structure annotation are as "
+            "The current values for each jaxtyping PyTree structure annotation are as "
             "follows."
         )
         for name, structure in pytree_memo.items():
             pieces.append(f"{name}={structure}")
     return "\n".join(pieces)
+
+
+class _jaxtyping_note_str(str):
+    """Used with `_no_jaxtyping_note` to flag that a note came from jaxtyping."""
+
+
+def _no_jaxtyping_note(e: Exception) -> bool:
+    """Checks if any of the exception's notes are from jaxtyping."""
+    try:
+        notes = e.__notes__
+    except AttributeError:
+        return True
+    else:
+        for note in notes:
+            if isinstance(note, _jaxtyping_note_str):
+                return False
+        return True
+
+
+_spacer = "--------------------\n"
diff --git a/jaxtyping/_import_hook.py b/jaxtyping/_import_hook.py
index ebe89e9..f2a551e 100644
--- a/jaxtyping/_import_hook.py
+++ b/jaxtyping/_import_hook.py
@@ -78,18 +78,13 @@ def _optimized_cache_from_source(typechecker_hash, /, path, debug_override=None)
     #    changing the typechecker will hit a different cache.
     # Version 7: Using the same md5 hash of the `typechecker` argument
     #    for importlib and decorator lookup.
+    # Version 8: Now using new-style `jaxtyped(typechecker=...)` rather than old-style
+    #    double-decorators.
     return cache_from_source(
-        path, debug_override, optimization=f"jaxtyping7{typechecker_hash}"
+        path, debug_override, optimization=f"jaxtyping8{typechecker_hash}"
     )
 
 
-def _dot_lookup(*elements):
-    out = ast.Name(id=elements[0], ctx=ast.Load())
-    for element in elements[1:]:
-        out = ast.Attribute(out, element, ctx=ast.Load())
-    return out
-
-
 class Typechecker:
     lookup = {}
 
@@ -130,7 +125,7 @@ class Typechecker:
         if self.ast is None:
             self.ast = (
                 ast.parse(
-                    f"@jaxtyping._import_hook.Typechecker.lookup['{self.hash}']\n"
+                    f"@jaxtyping.jaxtyped(typechecker=jaxtyping._import_hook.Typechecker.lookup['{self.hash}'])\n"
                     "def _():\n    ..."
                 )
                 .body[0]
@@ -162,10 +157,9 @@ class JaxtypingTransformer(ast.NodeVisitor):
         return node
 
     def visit_ClassDef(self, node: ast.ClassDef):
-        func = _dot_lookup("jaxtyping", "_decorator", "_jaxtyped_typechecker")
-        node.decorator_list.insert(
-            0, ast.Call(func, [self._typechecker.get_ast()], keywords=[])
-        )
+        # Place at the start of the decorator list, so that `@dataclass` decorators get
+        # called first.
+        node.decorator_list.insert(0, self._typechecker.get_ast())
         self._parents.append(node)
         self.generic_visit(node)
         self._parents.pop()
@@ -175,8 +169,13 @@ class JaxtypingTransformer(ast.NodeVisitor):
         has_annotated_args = any(arg for arg in node.args.args if arg.annotation)
         has_annotated_return = bool(node.returns)
         if has_annotated_args or has_annotated_return:
-            # Place at the end of the decorator list, as otherwise we wrap e.g.
-            # `jax.custom_{jvp,vjp}` and lose the ability to `defjvp` etc.
+            # Place at the end of the decorator list, because:
+            # - as otherwise we wrap e.g. `jax.custom_{jvp,vjp}` and lose the ability
+            #     to `defjvp` etc.
+            # - decorators frequently remove annotations from functions, and we'd like
+            #     to use those annotations.
+            # - typeguard in particular wants to be at the end of the decorator list, as
+            #     it works by recompling the wrapped function.
             #
             # Note that the counter-argument here is that we'd like to place this
             # at the start of the decorator list, in case a typechecking annotation
@@ -184,14 +183,6 @@ class JaxtypingTransformer(ast.NodeVisitor):
             # case we're just going to have to need to ask the user to remove their
             # typechecking annotation (and let this decorator do it instead).
             # It's more important we be compatible with normal JAX code.
-            #
-            #
-            # FWIW, typeguard also wants to be at the end of the decorator list, as it
-            # works by recompiling the wrapped function.
-            node.decorator_list.append(_dot_lookup("jaxtyping", "jaxtyped"))
-            # Place typechecker at the end of the decorator list, as decorators
-            # frequently remove annotations from functions and we'd like to
-            # use those annotations.
             node.decorator_list.append(self._typechecker.get_ast())
 
         self._parents.append(node)
@@ -290,7 +281,8 @@ class ImportHookManager:
 # Deliberately no default for `typechecker` so that folks must opt-in to not having
 # a typechecker.
 def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optional[str]):
-    """Automatically apply `@jaxtyped`, and optionally a type checker, as decorators.
+    """Automatically apply the `@jaxtyped(typechecker=typechecker)` decorator to every
+    function and dataclass over a whole codebase.
 
     !!! Tip "Usage"
 
@@ -298,18 +290,19 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
         from jaxtyping import install_import_hook
         # Plus any one of the following:
 
-        # decorate @jaxtyped and @typeguard.typechecked
+        # decorate `@jaxtyped(typechecker=typeguard.typechecked)`
         with install_import_hook("foo", "typeguard.typechecked"):
             import foo          # Any module imported inside this `with` block, whose
             import foo.bar      # name begins with the specified string, will
             import foo.bar.qux  # automatically have both `@jaxtyped` and the specified
-                                # typechecker applied to all of their functions.
+                                # typechecker applied to all of their functions and
+                                # dataclasses.
 
-        # decorate @jaxtyped and @beartype.beartype
+        # decorate `@jaxtyped(typechecker=beartype.beartype)`
         with install_import_hook("foo", "beartype.beartype"):
             ...
 
-        # decorate only @jaxtyped (if you want that for some reason)
+        # decorate only `@jaxtyped` (if you want that for some reason)
         with install_import_hook("foo", None):
             ...
         ```
@@ -326,22 +319,9 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
         install_import_hook(["foo", "bar.baz"], ...)
         ```
 
-    The import hook will automatically decorate all functions, and check the attributes
-    assigned to dataclasses.
-
-    If the function already has any decorators on it, then both the `@jaxtyped` and the
-    typechecker decorators will get added at the bottom of the decorator list, e.g.
-    ```python
-    @some_other_decorator
-    @jaxtyped
-    @beartype.beartype
-    def foo(...): ...
-    ```
-
     **Arguments:**:
 
-    - `modules`: the names of the modules in which to automatically apply `@jaxtyped`
-        and `@typechecked`.
+    - `modules`: the names of the modules in which to automatically apply `@jaxtyped`.
     - `typechecker`: the module and function of the typechecker you want to use, as a
         string. For example `typechecker="typeguard.typechecked"`, or
         `typechecker="beartype.beartype"`. You may pass `typechecker=None` if you do not
@@ -351,7 +331,7 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
 
     A context manager that uninstalls the hook on exit, or when you call `.uninstall()`.
 
-    ??? Example "Example: end-user script"
+    !!! Example "Example: end-user script"
 
         ```python
         ### entry_point.py
@@ -366,7 +346,7 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
             ...
         ```
 
-    ??? Example "Example: writing a library"
+    !!! Example "Example: writing a library"
 
         ```python
         ### __init__.py
@@ -378,30 +358,6 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
                                                  # so will be hook'd.
         ```
 
-    ??? info "Pytest hook"
-
-        The import hook can be installed at test-time only, as a pytest hook. From the
-        command line the syntax is:
-        ```
-        pytest --jaxtyping-packages=foo,bar.baz,beartype.beartype
-        ```
-        or in `pyproject.toml`:
-        ```toml
-        [tool.pytest.ini_options]
-        addopts = "--jaxtyping-packages=foo,bar.baz,beartype.beartype"
-        ```
-        or in `pytest.ini`:
-        ```ini
-        [pytest]
-        addopts = --jaxtyping-packages=foo,bar.baz,beartype.beartype
-        ```
-        This example will apply the import hook to all modules whose names start with
-        either `foo` or `bar.baz`. The typechecker used in this example is
-        `beartype.beartype`.
-
-        (This is the author's preferred approach to performing runtime type-checking
-        with jaxtyping!)
-
     !!! warning
 
         Stringified dataclass annotations, e.g.
@@ -423,6 +379,29 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
             x: tuple["int"]
         ```
         will likely raise an error, and must not be used at all.
+
+    !!! warning
+
+        If a function already has any decorators on it, then `@jaxtyped` will get added
+        at the bottom of the decorator list, e.g.
+        ```python
+        @some_other_decorator
+        @jaxtyped(typechecker=beartype.beartype)
+        def foo(...): ...
+        ```
+        This is to support the common case in which
+        `some_other_decorator = jax.custom_jvp` etc.
+
+        If a class already has any decorators in it, then `@jaxtyped` will get added to
+        the top of the decorator list, e.g.
+        ```python
+        @jaxtyped(typechecker=beartype.beartype)
+        @some_other_decorator
+        class A:
+            ...
+        ```
+        This is to support the common case in which
+        `some_other_decorator = dataclasses.dataclass`.
     """  # noqa: E501
 
     if isinstance(modules, str):
diff --git a/jaxtyping/_pytree_type.py b/jaxtyping/_pytree_type.py
index 10ebcdf..84ad9ec 100644
--- a/jaxtyping/_pytree_type.py
+++ b/jaxtyping/_pytree_type.py
@@ -19,11 +19,12 @@
 
 import functools as ft
 import typing
-from typing import Generic, TypeVar
+from typing import Any, Generic, TypeVar
 
 import jax.tree_util as jtu
 import typeguard
 
+from ._raise import jaxtyping_raise_from
 from ._storage import (
     clear_treepath_memo,
     get_shape_memo,
@@ -68,30 +69,41 @@ class _MetaPyTree(type):
             return False
 
     def _check(cls, obj, pytree_memo):
-        # We could use `isinstance` here but that would fail for more complicated
-        # types, e.g. PyTree[tuple[int]]. So at least internally we make a particular
-        # choice of typechecker.
-        #
-        # Deliberately not using @jaxtyped so that we share the same `memo` as whatever
-        # dynamic context we're currently in.
-        @typeguard.typechecked
-        def accepts_leaftype(x: cls.leaftype):
-            pass
+        if cls.leaftype is Any:
 
-        def is_leaftype(x, new_scope=True):
-            if new_scope and cls.structure is not None:
-                set_treepath_memo(None, cls.structure)
-            try:
-                accepts_leaftype(x)
-            except _TypeCheckError:
+            def is_flatten_leaftype(x):
                 return False
-            else:
-                return True
-            finally:
-                if new_scope and cls.structure is not None:
-                    clear_treepath_memo()
 
-        leaves, structure = jtu.tree_flatten(obj, is_leaf=is_leaftype)
+            def is_check_leaftype(x, new_scope):
+                return True
+
+        else:
+            # We could use `isinstance` here but that would fail for more complicated
+            # types, e.g. PyTree[tuple[int]]. So at least internally we make a
+            # particular choice of typechecker.
+            #
+            # Deliberately not using @jaxtyped so that we share the same `memo` as
+            # whatever dynamic context we're currently in.
+            @typeguard.typechecked
+            def accepts_leaftype(x: cls.leaftype):
+                pass
+
+            def is_leaftype(x, new_scope=True):
+                if new_scope and cls.structure is not None:
+                    set_treepath_memo(None, cls.structure)
+                try:
+                    accepts_leaftype(x)
+                except _TypeCheckError:
+                    return False
+                else:
+                    return True
+                finally:
+                    if new_scope and cls.structure is not None:
+                        clear_treepath_memo()
+
+            is_flatten_leaftype = is_check_leaftype = is_leaftype
+
+        leaves, structure = jtu.tree_flatten(obj, is_leaf=is_flatten_leaftype)
         if cls.structure is not None:
             if cls.structure.isidentifier():
                 try:
@@ -119,10 +131,14 @@ class _MetaPyTree(type):
                     try:
                         prev_structure = pytree_memo[identifier]
                     except KeyError as e:
-                        raise NameError(
-                            f"Cannot process composite structure '{cls.structure}' as "
-                            f"the structure name {identifier} has not been seen before."
-                        ) from e
+                        jaxtyping_raise_from(
+                            NameError(
+                                f"Cannot process composite structure '{cls.structure}' "
+                                f"as the structure name {identifier} has not been seen "
+                                "before."
+                            ),
+                            e,
+                        )
                     # Not using `PyTreeDef.compose` due to JAX bug #18218.
                     prev_pytree = jtu.tree_unflatten(
                         prev_structure, [0] * prev_structure.num_leaves
@@ -152,7 +168,7 @@ class _MetaPyTree(type):
             for leaf_index, leaf in enumerate(leaves):
                 if cls.structure is not None:
                     set_treepath_memo(leaf_index, cls.structure)
-                if not is_leaftype(leaf, new_scope=False):
+                if not is_check_leaftype(leaf, new_scope=False):
                     return False
                 clear_treepath_memo()
         finally:
diff --git a/jaxtyping/_raise.py b/jaxtyping/_raise.py
new file mode 100644
index 0000000..9230877
--- /dev/null
+++ b/jaxtyping/_raise.py
@@ -0,0 +1,22 @@
+from typing import NoReturn
+
+
+def jaxtyping_raise(e) -> NoReturn:
+    """Raises `e`, whilst adding a tag that it should not be intercepted by
+    `TypeCheckError`. All `raise` statements from within `__instancecheck__` should use
+    this.
+    """
+    try:
+        raise e
+    except Exception as f:
+        f._jaxtyping_malformed = True
+        raise
+
+
+def jaxtyping_raise_from(e, e_base) -> NoReturn:
+    __tracebackhide__ = True
+    try:
+        raise e from e_base
+    except Exception as f:
+        f._jaxtyping_malformed = True
+        raise
diff --git a/jaxtyping/_storage.py b/jaxtyping/_storage.py
index 648d70a..ed9b3bc 100644
--- a/jaxtyping/_storage.py
+++ b/jaxtyping/_storage.py
@@ -20,6 +20,8 @@
 import threading
 from typing import Optional
 
+from ._raise import jaxtyping_raise
+
 
 _shape_storage = threading.local()
 
@@ -70,10 +72,12 @@ def clear_treepath_memo() -> None:
 
 def set_treepath_memo(index: Optional[int], structure: str) -> None:
     if hasattr(_treepath_storage, "value") and _treepath_storage.value is not None:
-        raise ValueError(
-            "Cannot typecheck annotations of the form "
-            "`PyTree[PyTree[Shaped[Array, '?foo'], 'T'], 'S']` as it is ambiguous "
-            "which PyTree the `?` annotation refers to."
+        jaxtyping_raise(
+            ValueError(
+                "Cannot typecheck annotations of the form "
+                "`PyTree[PyTree[Shaped[Array, '?foo'], 'T'], 'S']` as it is ambiguous "
+                "which PyTree the `?` annotation refers to."
+            )
         )
     if index is None:
         _treepath_storage.value = f"~~delete~~({structure}) "
@@ -84,9 +88,11 @@ def set_treepath_memo(index: Optional[int], structure: str) -> None:
 
 def get_treepath_memo() -> str:
     if not hasattr(_treepath_storage, "value") or _treepath_storage.value is None:
-        raise ValueError(
-            "Cannot use `?` annotations, e.g. `Shaped[Array, '?foo']`, except "
-            "when contained with structured `PyTree` annotations, e.g. "
-            "`PyTree[Shaped[Array, '?foo'], 'T']`."
+        jaxtyping_raise(
+            ValueError(
+                "Cannot use `?` annotations, e.g. `Shaped[Array, '?foo']`, except "
+                "when contained with structured `PyTree` annotations, e.g. "
+                "`PyTree[Shaped[Array, '?foo'], 'T']`."
+            )
         )
     return _treepath_storage.value
diff --git a/test/conftest.py b/test/conftest.py
index ca4e5d4..896e4cb 100644
--- a/test/conftest.py
+++ b/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():
diff --git a/test/test_array.py b/test/test_array.py
index 77798cf..4f3fc26 100644
--- a/test/test_array.py
+++ b/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):
         ...
 
diff --git a/test/test_decorator.py b/test/test_decorator.py
index f64bf57..331a3d5 100644
--- a/test/test_decorator.py
+++ b/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.
diff --git a/test/test_messages.py b/test/test_messages.py
new file mode 100644
index 0000000..2d8f0ad
--- /dev/null
+++ b/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)
diff --git a/test/test_pytree.py b/test/test_pytree.py
index 169f05c..a1b416e 100644
--- a/test/test_pytree.py
+++ b/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