From 59e8fb0d18325f990a9d59ee35e90c04b699cab8 Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Fri, 30 Dec 2022 19:00:26 +0000
Subject: [PATCH] Hopefully fixed PyTree raising spurious errors. Bit
 mysterious that this worked before, really. I've tested this fix as best I
 can against the various static type checkers, but these are weird and varied
 enough that this might not be a perfect fix. If you see this and have issues,
 let me know. (#54)

---
 jaxtyping/__init__.py | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
index 74f65d1..42ec0e7 100644
--- a/jaxtyping/__init__.py
+++ b/jaxtyping/__init__.py
@@ -79,14 +79,27 @@ from .import_hook import install_import_hook as install_import_hook
 
 
 if typing.TYPE_CHECKING:
-    _T = typing.TypeVar("_T")
-
-    class PyTree(typing_extensions.Protocol[_T]):
-        pass
-
+    # Set up to deliberately confuse a static type checker.
+    PyTree = getattr(typing, "foo" + "bar")
+    # What's going on with this madness?
+    #
+    # At static-type-checking-time, we want `PyTree` to be a type for which both
+    # `PyTree` and `PyTree[Foo]` are equivalent to `Any`.
+    # (The intention is that `PyTree` be a runtime-only type; there's no real way to
+    # do more with static type checkers.)
+    #
+    # Unfortunately, this isn't possible: `Any` isn't subscriptable. And there's no
+    # equivalent way we can fake this using typing annotations. (In some sense the
+    # closest thing would be a `Protocol[T]` with no methods, but that's actually the
+    # opposite of what we want: that ends up allowing nothing at all.)
+    #
+    # The good news for us is that static type checkers have an internal escape hatch.
+    # If they can't figure out what a type is, then they just give up and allow
+    # anything. (I believe this is sometimes called `Unknown`.) Thus, this odd-looking
+    # annotation, which static type checkers aren't smart enough to resolve.
 elif has_jax:
     from .pytree_type import PyTree
 
 del has_jax
 
-__version__ = "0.2.10"
+__version__ = "0.2.11"