Fixes for some new failures. (Where did they come from?) (#65)

* Fixes for some new failures. (Where did they come from?)

* Fixed isort?

.github/workflows/run_tests.yml

@@ -26,7 +26,7 @@ jobs:
   run-tests:
     strategy:
       matrix:
-        python-version: [ 3.7, 3.8, 3.9 ]
+        python-version: [ 3.8, 3.9 ]
         os: [ ubuntu-latest ]
       fail-fast: false
     runs-on: ${{ matrix.os }}

.pre-commit-config.yaml

@@ -23,13 +23,13 @@ repos:
     hooks:
     - id: black
   - repo: https://github.com/nbQA-dev/nbQA
-    rev: 1.2.3
+    rev: 1.6.3
     hooks:
     - id: nbqa-black
     - id: nbqa-isort 
     - id: nbqa-flake8
   - repo: https://github.com/PyCQA/isort
-    rev: 5.10.1
+    rev: 5.12.0
     hooks:
     - id: isort
   - repo: https://github.com/pycqa/flake8

README.md

@@ -2,9 +2,10 @@
 
 Type annotations **and runtime checking** for:
 
-1. shape and dtype of [JAX](https://github.com/google/jax) arrays;
+1. shape and dtype of [JAX](https://github.com/google/jax) arrays; *(Now also supports PyTorch, NumPy, and TensorFlow!)*
 2. [PyTrees](https://jax.readthedocs.io/en/latest/pytrees.html).
 
+
 **For example:**
 ```python
 from jaxtyping import Array, Float, PyTree
@@ -28,8 +29,6 @@ def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
 pip install jaxtyping
 ```
 
-Requires JAX 0.3.4+.
-
 Also install your favourite runtime type-checking package. The two most popular are [typeguard](https://github.com/agronholm/typeguard) (which exhaustively checks every argument) and [beartype](https://github.com/beartype/beartype) (which checks random pieces of arguments).
 
 ## Documentation

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.9"
+__version__ = "0.2.12"

jaxtyping/import_hook.py

@@ -66,9 +66,11 @@ def _call_with_frames_removed(f, *args, **kwargs):
 
 def _optimized_cache_from_source(path, debug_override=None):
     # Version 2: change the position of the `@jaxtyped` decorator, so need a
-    # different name to avoid hitting old __pycache__
+    #     different name to avoid hitting old __pycache__.
     # Version 3: now also annotating classes.
-    return cache_from_source(path, debug_override, optimization="jaxtyping3")
+    # Version 4: I'm honestly not sure, but bumping this fixed some kind of odd error.
+    #     Maybe I changed something with hte classes part way through version 3?
+    return cache_from_source(path, debug_override, optimization="jaxtyping4")
 
 
 def _dot_lookup(*elements):
@@ -111,7 +113,7 @@ class _JaxtypingTransformer(ast.NodeVisitor):
             args = [ast.Constant(None)]
         else:
             args = [_dot_lookup(*self._typechecker)]
-        node.decorator_list.append(ast.Call(func, args, keywords=[]))
+        node.decorator_list.insert(0, ast.Call(func, args, keywords=[]))
         self._parents.append(node)
         self.generic_visit(node)
         self._parents.pop()

jaxtyping/pytree_type.py

@@ -59,6 +59,14 @@ class _MetaPyTree(type):
         return out
 
 
+try:
+    # new typeguard
+    _TypeCheckError = (TypeError, typeguard.TypeCheckError)
+except AttributeError:
+    # old typeguard
+    _TypeCheckError = TypeError
+
+
 class _MetaSubscriptPyTree(type):
     def __call__(self, *args, **kwargs):
         raise RuntimeError("PyTree cannot be instantiated")
@@ -77,7 +85,7 @@ class _MetaSubscriptPyTree(type):
         def is_leaftype(x):
             try:
                 accepts_leaftype(x)
-            except TypeError:
+            except _TypeCheckError:
                 return False
             else:
                 return True

test/helpers.py

@@ -18,16 +18,31 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import equinox as eqx
+import typeguard
 
 
+ParamError = []
+ReturnError = []
+ParamError.append(TypeError)  # old typeguard
+ReturnError.append(TypeError)  # old typeguard
+
+try:
+    # new typeguard
+    ParamError.append(typeguard.TypeCheckError)
+    ReturnError.append(typeguard.TypeCheckError)
+except AttributeError:
+    pass
+
 try:
     import beartype
 except ImportError:
-    ParamError = TypeError
-    ReturnError = TypeError
+    pass
 else:
-    ParamError = (TypeError, beartype.roar.BeartypeCallHintParamViolation)
-    ReturnError = (TypeError, beartype.roar.BeartypeCallHintReturnViolation)
+    ParamError.append(beartype.roar.BeartypeCallHintParamViolation)
+    ReturnError.append(beartype.roar.BeartypeCallHintReturnViolation)
+
+ParamError = tuple(ParamError)
+ReturnError = tuple(ReturnError)
 
 
 @eqx.filter_jit

test/import_hook_tester_beartype.py

@@ -17,6 +17,8 @@
 # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+import dataclasses
+
 import equinox as eqx
 import jax.numpy as jnp
 import pytest
@@ -45,3 +47,16 @@ with pytest.raises(ParamError):
     M(1.0, jnp.array([1.0]))
 with pytest.raises(ParamError):
     M(1, jnp.array(1.0))
+
+
+@dataclasses.dataclass
+class D:
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+D(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1, jnp.array(1.0))

test/import_hook_tester_broken_checker.py

@@ -17,6 +17,8 @@
 # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+import dataclasses
+
 import equinox as eqx
 import jax.numpy as jnp
 import pytest
@@ -45,3 +47,16 @@ with pytest.raises(ParamError):
     M(1.0, jnp.array([1.0]))
 with pytest.raises(ParamError):
     M(1, jnp.array(1.0))
+
+
+@dataclasses.dataclass
+class D:
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+D(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1, jnp.array(1.0))

test/import_hook_tester_typeguard.py

@@ -17,6 +17,8 @@
 # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+import dataclasses
+
 import equinox as eqx
 import jax.numpy as jnp
 import pytest
@@ -45,3 +47,16 @@ with pytest.raises(ParamError):
     M(1.0, jnp.array([1.0]))
 with pytest.raises(ParamError):
     M(1, jnp.array(1.0))
+
+
+@dataclasses.dataclass
+class D:
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+D(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1, jnp.array(1.0))

test/test_decorator.py

@@ -1,13 +1,26 @@
+import abc
+
 from jaxtyping import jaxtyped
 
 
-class M:
+class M(metaclass=abc.ABCMeta):
     @jaxtyped
     @classmethod
     def f(cls):
         return 3
 
+    @jaxtyped
+    @abc.abstractmethod
+    def g(self):
+        ...
+
 
 # Check that the @jaxtyped decorator doesn't blat the __get__ of @classmethod
-def test_decorator():
+def test_classmethod():
     assert M.f() == 3
+
+
+# Check that the @jaxtyped decorator doesn't blat the __isabstractmethod__ of
+# @abstractmethod
+def test_abstractmethod():
+    assert M.g.__isabstractmethod__