Error message improvements

jaxtyping/_decorator.py

@@ -687,17 +687,25 @@ def _exc_shape_info(memos) -> str:
     Used in type-checking error messages.
     """
     single_memo, variadic_memo, pytree_memo, _ = memos
+    single_memo = {
+        name: size
+        for name, size in single_memo.items()
+        if not name.startswith("~~delete~~")
+    }
+    variadic_memo = {
+        name: shape
+        for name, (_, shape) in variadic_memo.items()
+        if not name.startswith("~~delete~~")
+    }
     pieces = []
     if len(single_memo) > 0 or len(variadic_memo) > 0:
         pieces.append(
             "The current values for each jaxtyping axis annotation are as follows."
         )
         for name, size in single_memo.items():
-            if not name.startswith("~~delete~~"):
-                pieces.append(f"{name}={size}")
-        for name, (_, shape) in variadic_memo.items():
-            if not name.startswith("~~delete~~"):
-                pieces.append(f"{name}={shape}")
+            pieces.append(f"{name}={size}")
+        for name, shape in variadic_memo.items():
+            pieces.append(f"{name}={shape}")
     if len(pytree_memo) > 0:
         pieces.append(
             "The current values for each jaxtyping PyTree structure annotation are as "

jaxtyping/_raise.py

@@ -6,6 +6,7 @@ def jaxtyping_raise(e) -> NoReturn:
     `TypeCheckError`. All `raise` statements from within `__instancecheck__` should use
     this.
     """
+    __tracebackhide__ = True
     try:
         raise e
     except Exception as f:

test/test_decorator.py

@@ -110,7 +110,7 @@ def test_varkwargs(jaxtyp, typecheck):
 
 def test_defaults(jaxtyp, typecheck):
     @jaxtyp(typecheck)
-    def f(x, y=1):
+    def f(x: int, y=1):
         pass
 
     f(1)