jaxtyping v0.2.29 now wraps functions without annotations as well. That changes some tracebacks slightly.

equinox/_module.py

@@ -1061,6 +1061,12 @@ def __wrapped__(self):
             self.__self__, type(self.__self__)
         )
 
+    # This should be unnecessary in principle. In practice something goes wrong on
+    # Python 3.9 and it returns the wrong thing.
+    @property
+    def __signature__(self):
+        return inspect.signature(self.__wrapped__)
+
 
 #
 # Part 3: some downstream pieces. These don't actually affect the core `Module`

tests/test_errors.py

@@ -146,11 +146,15 @@ def g(x):
         tb = e.__traceback__
         code_stack = []
         while tb is not None:
-            code_stack.append(tb.tb_frame.f_code)
+            if not tb.tb_frame.f_globals["__name__"].startswith("jaxtyping"):
+                code_stack.append(tb.tb_frame.f_code)
             tb = tb.tb_next
-        assert len(code_stack) == 1
-        assert code_stack[0].co_filename.endswith("test_errors.py")
-        assert code_stack[0].co_name == "test_traceback_runtime_eqx"
+        assert len(code_stack) == 2
+        one, two = code_stack
+        assert one.co_filename.endswith("test_errors.py")
+        assert one.co_name == "test_traceback_runtime_eqx"
+        assert two.co_filename.endswith("equinox/_jit.py")
+        assert two.co_name == "_call"
 
 
 def test_traceback_runtime_custom():
@@ -177,7 +181,8 @@ def _raises():
         tb = e.__traceback__
         code_stack = []
         while tb is not None:
-            code_stack.append(tb.tb_frame.f_code)
+            if not tb.tb_frame.f_globals["__name__"].startswith("jaxtyping"):
+                code_stack.append(tb.tb_frame.f_code)
             tb = tb.tb_next
         assert len(code_stack) == 4
         one, two, three, four = code_stack