diff --git a/API.md b/API.md
index c1ad912..104cfa6 100644
--- a/API.md
+++ b/API.md
@@ -168,6 +168,8 @@ The import hook can be applied to multiple packages via
 install_import_hook(["foo", "bar.baz"], ...)
 ```
 
+The import hook will automatically decorate all functions, and the `__init__` method of dataclasses.
+
 **Example: writing an end-user script**
 
 ```python
diff --git a/jaxtyping/decorator.py b/jaxtyping/decorator.py
index 61d82bb..a269d00 100644
--- a/jaxtyping/decorator.py
+++ b/jaxtyping/decorator.py
@@ -17,6 +17,7 @@
 # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+import dataclasses
 import functools as ft
 import inspect
 import threading
@@ -46,8 +47,30 @@ class _Jaxtyped:
 
 def jaxtyped(fn):
     if inspect.isclass(fn):  # allow decorators on class definitions
-        init = jaxtyped(fn.__init__)
-        fn.__init__ = init
-        return fn
+        if dataclasses.is_dataclass(fn):
+            init = jaxtyped(fn.__init__)
+            fn.__init__ = init
+            return fn
+        else:
+            raise ValueError(
+                "jaxtyped may only be added as a class decorator to dataclasses"
+            )
     else:
         return ft.wraps(fn)(_Jaxtyped(fn))
+
+
+def _jaxtyped_typechecker(typechecker):
+    # typechecker is expected to probably be either `typeguard.typechecked`, or
+    # `beartype.beartype`, or `None`.
+
+    if typechecker is None:
+        typechecker = lambda x: x
+
+    def _wrapper(kls):
+        assert inspect.isclass(kls)
+        if dataclasses.is_dataclass(kls):
+            init = jaxtyped(typechecker(kls.__init__))
+            kls.__init__ = init
+        return kls
+
+    return _wrapper
diff --git a/jaxtyping/import_hook.py b/jaxtyping/import_hook.py
index a324bf3..6d66e76 100644
--- a/jaxtyping/import_hook.py
+++ b/jaxtyping/import_hook.py
@@ -67,7 +67,15 @@ def _call_with_frames_removed(f, *args, **kwargs):
 def _optimized_cache_from_source(path, debug_override=None):
     # Version 2: change the position of the `@jaxtyped` decorator, so need a
     # different name to avoid hitting old __pycache__
-    return cache_from_source(path, debug_override, optimization="jaxtyping2")
+    # Version 3: now also annotating classes.
+    return cache_from_source(path, debug_override, optimization="jaxtyping3")
+
+
+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 _JaxtypingTransformer(ast.NodeVisitor):
@@ -97,6 +105,18 @@ class _JaxtypingTransformer(ast.NodeVisitor):
         self._parents.pop()
         return node
 
+    def visit_ClassDef(self, node: ast.ClassDef):
+        func = _dot_lookup("jaxtyping", "decorator", "_jaxtyped_typechecker")
+        if self._typechecker is None:
+            args = [ast.Constant(None)]
+        else:
+            args = [_dot_lookup(*self._typechecker)]
+        node.decorator_list.append(ast.Call(func, args, keywords=[]))
+        self._parents.append(node)
+        self.generic_visit(node)
+        self._parents.pop()
+        return node
+
     def visit_FunctionDef(self, node: ast.FunctionDef):
         has_annotated_args = any(arg for arg in node.args.args if arg.annotation)
         has_annotated_return = bool(node.returns)
@@ -110,24 +130,12 @@ 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.
-            node.decorator_list.append(
-                ast.Attribute(
-                    ast.Name(id="jaxtyping", ctx=ast.Load()), "jaxtyped", ast.Load()
-                ),
-            )
+            node.decorator_list.append(_dot_lookup("jaxtyping", "jaxtyped"))
             if self._typechecker is not None:
                 # Place at the end of the decorator list, as decorators
                 # frequently remove annotations from functions and we'd like to
                 # use those annotations.
-                typechecker_module, typechecker_function = self._typechecker
-                node.decorator_list.append(
-                    ast.Attribute(
-                        ast.Name(id=typechecker_module, ctx=ast.Load()),
-                        typechecker_function,
-                        ast.Load(),
-                    )
-                )
-
+                node.decorator_list.append(_dot_lookup(*self._typechecker))
         self._parents.append(node)
         self.generic_visit(node)
         self._parents.pop()
diff --git a/test/import_hook_tester_beartype.py b/test/import_hook_tester_beartype.py
index 414fa47..db1cd29 100644
--- a/test/import_hook_tester_beartype.py
+++ b/test/import_hook_tester_beartype.py
@@ -17,6 +17,7 @@
 # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+import equinox as eqx
 import jax.numpy as jnp
 import pytest
 
@@ -32,3 +33,15 @@ def g(x: Float32[jnp.ndarray, " b"]):
 g(jnp.array([1.0]))
 with pytest.raises(ParamError):
     g(jnp.array(1))
+
+
+class M(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+M(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1, jnp.array(1.0))
diff --git a/test/import_hook_tester_broken_checker.py b/test/import_hook_tester_broken_checker.py
index 414fa47..db1cd29 100644
--- a/test/import_hook_tester_broken_checker.py
+++ b/test/import_hook_tester_broken_checker.py
@@ -17,6 +17,7 @@
 # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+import equinox as eqx
 import jax.numpy as jnp
 import pytest
 
@@ -32,3 +33,15 @@ def g(x: Float32[jnp.ndarray, " b"]):
 g(jnp.array([1.0]))
 with pytest.raises(ParamError):
     g(jnp.array(1))
+
+
+class M(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+M(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1, jnp.array(1.0))
diff --git a/test/import_hook_tester_transitive/another_file.py b/test/import_hook_tester_transitive/another_file.py
index a7f710f..a9cfb0b 100644
--- a/test/import_hook_tester_transitive/another_file.py
+++ b/test/import_hook_tester_transitive/another_file.py
@@ -17,6 +17,7 @@
 # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+import equinox as eqx
 import jax.numpy as jnp
 import pytest
 
@@ -32,3 +33,15 @@ def g(x: Float32[jnp.ndarray, " b"]):
 g(jnp.array([1.0]))
 with pytest.raises(ParamError):
     g(jnp.array(1))
+
+
+class M(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+M(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1, jnp.array(1.0))
diff --git a/test/import_hook_tester_typeguard.py b/test/import_hook_tester_typeguard.py
index 414fa47..db1cd29 100644
--- a/test/import_hook_tester_typeguard.py
+++ b/test/import_hook_tester_typeguard.py
@@ -17,6 +17,7 @@
 # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+import equinox as eqx
 import jax.numpy as jnp
 import pytest
 
@@ -32,3 +33,15 @@ def g(x: Float32[jnp.ndarray, " b"]):
 g(jnp.array([1.0]))
 with pytest.raises(ParamError):
     g(jnp.array(1))
+
+
+class M(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+M(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1, jnp.array(1.0))
diff --git a/test/test_import_hook.py b/test/test_import_hook.py
index 9077082..66c6cc6 100644
--- a/test/test_import_hook.py
+++ b/test/test_import_hook.py
@@ -26,9 +26,8 @@ def test_import_hook_typeguard():
     hook = install_import_hook(
         "test.import_hook_tester_typeguard", ("typeguard", "typechecked")
     )
-    from . import import_hook_tester_typeguard  # noqa: F401
-
-    hook.uninstall()
+    with hook:
+        from . import import_hook_tester_typeguard  # noqa: F401
 
 
 def test_import_hook_beartype():
@@ -40,24 +39,21 @@ def test_import_hook_beartype():
         hook = install_import_hook(
             "test.import_hook_tester_beartype", ("beartype", "beartype")
         )
-        from . import import_hook_tester_beartype  # noqa: F401
-
-        hook.uninstall()
+        with hook:
+            from . import import_hook_tester_beartype  # noqa: F401
 
 
 def test_import_hook_transitive():
     hook = install_import_hook(
         "test.import_hook_tester_transitive", ("typeguard", "typechecked")
     )
-    from . import import_hook_tester_transitive  # noqa: F401
-
-    hook.uninstall()
+    with hook:
+        from . import import_hook_tester_transitive  # noqa: F401
 
 
 def test_import_hook_broken_checker():
     hook = install_import_hook(
         "test.import_hook_tester_broken_checker", ("jaxtyping", "does_not_exist")
     )
-    with pytest.raises(AttributeError):
+    with hook, pytest.raises(AttributeError):
         from . import import_hook_tester_broken_checker  # noqa: F401
-    hook.uninstall()