From 4c90808401493b705b595f477cad77fe22bab081 Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Wed, 26 Jul 2023 09:26:59 -0700
Subject: [PATCH] jaxtyping no longer appears in tracebacks

---
 jaxtyping/_decorator.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py
index 07abcfd..a0c4f13 100644
--- a/jaxtyping/_decorator.py
+++ b/jaxtyping/_decorator.py
@@ -25,6 +25,14 @@ import types
 import weakref
 
 
+try:
+    import jax._src.traceback_util as traceback_util
+except ImportError:
+    pass
+else:
+    traceback_util.register_exclusion(__file__)
+
+
 storage = threading.local()