From d6bf88c066f3d8be3bf918e970e6154f19417f10 Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Fri, 31 May 2024 09:26:30 +0200
Subject: [PATCH 1/2] jaxtyping v0.2.29 now wraps functions without annotations
 as well. That changes some tracebacks slightly.

---
 tests/test_errors.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/tests/test_errors.py b/tests/test_errors.py
index c39f9950..2a9794cd 100644
--- a/tests/test_errors.py
+++ b/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

From 1def460111af38191d00f998e2a9047e6cce2a69 Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Fri, 7 Jun 2024 22:12:23 +0200
Subject: [PATCH 2/2] Fixed spurious Python 3.9 error.

---
 equinox/_module.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/equinox/_module.py b/equinox/_module.py
index fe37b72f..da8a8b49 100644
--- a/equinox/_module.py
+++ b/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`