Bump version

* 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

API.md

@@ -62,6 +62,8 @@ Float32[Array, "some_shape"]
 
 The array should typically be a `jaxtyping.Array`, which is an alias for `jax.numpy.ndarray`.
 
+`jaxtyping.ArrayLike` is also available, which is an alias for `jax.typing.ArrayLike`. This is a union over JAX arrays and the builtin `bool`/`int`/`float`/`complex`.
+
 But you can use other types as well. `jaxtyping` has support for JAX, NumPy, TensorFlow, and PyTorch, e.g.:
 ```python
 Float[np.ndarray, "..."]

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,6 +29,10 @@ def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
 pip install jaxtyping
 ```
 
+Requires Python 3.8+.
+
+JAX is an optional dependency, required for `jaxtyping.{Array, ArrayLike, PyTree}`. If JAX is not installed then these types will not be available, but you may still use jaxtyping alongside PyTorch/NumPy/etc.
+
 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
@@ -48,12 +53,6 @@ Computer vision models: [Eqxvision](https://github.com/paganpasta/eqxvision).
 
 SymPy<->JAX conversion; train symbolic expressions via gradient descent: [sympy2jax](https://github.com/google/sympy2jax).
 
-### Acknowledgements
-
-Shape annotations + runtime type checking is inspired by [TorchTyping](https://github.com/patrick-kidger/torchtyping).
-
-The concise syntax is partially inspired by [etils.array_types](https://github.com/google/etils/tree/main/etils/array_types).
-
 ### Disclaimer
 
 This is not an official Google product.

jaxtyping/__init__.py

@@ -18,7 +18,6 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import typing
-import typing_extensions
 
 
 try:
@@ -35,6 +34,7 @@ if typing.TYPE_CHECKING:
     # For imports, we need to explicitly `import X as X` in order for Pyright to see
     # them as public. See discussion at https://github.com/microsoft/pyright/issues/2277
     from jax import Array as Array
+    from jax.typing import ArrayLike as ArrayLike
 elif has_jax:
     if getattr(typing, "GENERATING_DOCUMENTATION", False):
 
@@ -42,9 +42,19 @@ elif has_jax:
             pass
 
         Array.__module__ = "builtins"
+
+        class ArrayLike:
+            pass
+
+        ArrayLike.__module__ = "builtins"
     else:
         from jax import Array as Array
 
+        try:
+            from jax.typing import ArrayLike as ArrayLike
+        except (ModuleNotFoundError, ImportError):
+            pass
+
 from .array_types import (
     AbstractArray as AbstractArray,
     AbstractDtype as AbstractDtype,
@@ -79,14 +89,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.13"

jaxtyping/array_types.py

@@ -20,8 +20,17 @@
 import enum
 import functools as ft
 import typing
-from typing import Any, Dict, List, NoReturn, Optional, Tuple, TYPE_CHECKING, Union
-from typing_extensions import Literal
+from typing import (
+    Any,
+    Dict,
+    List,
+    Literal,
+    NoReturn,
+    Optional,
+    Tuple,
+    TYPE_CHECKING,
+    Union,
+)
 
 import numpy as np
 
@@ -58,24 +67,72 @@ class _NamedDim:
         self.name = name
         self.broadcastable = broadcastable
 
+    def __eq__(self, other):
+        if type(self) is not type(other):
+            return False
+        if self.name != other.name:
+            return False
+        if self.broadcastable != other.broadcastable:
+            return False
+        return True
+
+    def __hash__(self):
+        return hash((self.name, self.broadcastable))
+
 
 class _NamedVariadicDim:
     def __init__(self, name, broadcastable):
         self.name = name
         self.broadcastable = broadcastable
 
+    def __eq__(self, other):
+        if type(self) is not type(other):
+            return False
+        if self.name != other.name:
+            return False
+        if self.broadcastable != other.broadcastable:
+            return False
+        return True
+
+    def __hash__(self):
+        return hash((self.name, self.broadcastable))
+
 
 class _FixedDim:
     def __init__(self, size, broadcastable):
         self.size = size
         self.broadcastable = broadcastable
 
+    def __eq__(self, other):
+        if type(self) is not type(other):
+            return False
+        if self.size != other.size:
+            return False
+        if self.broadcastable != other.broadcastable:
+            return False
+        return True
+
+    def __hash__(self):
+        return hash((self.size, self.broadcastable))
+
 
 class _SymbolicDim:
     def __init__(self, expr, broadcastable):
         self.expr = expr
         self.broadcastable = broadcastable
 
+    def __eq__(self, other):
+        if type(self) is not type(other):
+            return False
+        if self.expr != other.expr:
+            return False
+        if self.broadcastable != other.broadcastable:
+            return False
+        return True
+
+    def __hash__(self):
+        return hash((self.expr, self.broadcastable))
+
 
 _AbstractDimOrVariadicDim = Union[
     Literal[_anonymous_dim],
@@ -127,6 +184,22 @@ def _check_dims(
 
 
 class _MetaAbstractArray(type):
+    def __eq__(self, other):
+        if type(self) is not type(other):
+            return False
+        if self.array_type is not other.array_type:
+            return False
+        if self.dtypes != other.dtypes:
+            return False
+        if self.dims != other.dims:
+            return False
+        if self.index_variadic != other.index_variadic:
+            return False
+        return True
+
+    def __hash__(self):
+        return hash((self.array_type, self.dtypes, self.dims, self.index_variadic))
+
     def __instancecheck__(cls, obj):
         if not isinstance(obj, cls.array_type):
             return False
@@ -234,6 +307,12 @@ class _MetaAbstractArray(type):
             assert False
 
 
+def _check_scalar(dtype, dtypes, dims):
+    if len(dims) != 0:
+        return False
+    return (_any_dtype is dtypes) or any(d.startswith(dtype) for d in dtypes)
+
+
 class AbstractArray(metaclass=_MetaAbstractArray):
     array_type: Any
     dtypes: List[str]
@@ -383,31 +462,68 @@ class _MetaAbstractDtype(type):
                 elem = compile(elem, "<string>", "eval")
                 elem = _SymbolicDim(elem, broadcastable)
             dims.append(elem)
-        # In python 3.8, e.g., typing.Union lacks `__name__`.
-        try:
-            type_str = array_type.__name__
-        except AttributeError:
-            type_str = repr(array_type)
-        if _array_name_format == "dtype_and_shape":
-            name = f"{cls.__name__}[{type_str}, '{dim_str}']"
-        elif _array_name_format == "array":
-            name = type_str
+        dims = tuple(dims)
+
+        _not_made = object()
+
+        def _make(x):
+            # Allow Python built-in numeric types.
+            # TODO: do something more generic than this? Should we _make all types
+            # that have `shape` and `dtype` attributes or something?
+            if x is bool:
+                if _check_scalar("bool", cls.dtypes, dims):
+                    return x
+                else:
+                    return _not_made
+            elif x is int:
+                if _check_scalar("int", cls.dtypes, dims):
+                    return x
+                else:
+                    return _not_made
+            elif x is float:
+                if _check_scalar("float", cls.dtypes, dims):
+                    return x
+                else:
+                    return _not_made
+            elif x is complex:
+                if _check_scalar("complex", cls.dtypes, dims):
+                    return x
+                else:
+                    return _not_made
+            try:
+                type_str = x.__name__
+            except AttributeError:
+                type_str = repr(x)
+            if _array_name_format == "dtype_and_shape":
+                name = f"{cls.__name__}[{type_str}, '{dim_str}']"
+            elif _array_name_format == "array":
+                name = type_str
+            else:
+                raise ValueError(
+                    f"array_name_format {_array_name_format} not recognised"
+                )
+            out = _MetaAbstractArray(
+                name,
+                (AbstractArray,),
+                dict(
+                    array_type=x,
+                    dtypes=cls.dtypes,
+                    dims=dims,
+                    index_variadic=index_variadic,
+                ),
+            )
+            if getattr(typing, "GENERATING_DOCUMENTATION", False):
+                out.__module__ = "builtins"
+            else:
+                out.__module__ = "jaxtyping"
+            return out
+
+        if typing.get_origin(array_type) is typing.Union:
+            out = [_make(x) for x in typing.get_args(array_type)]
+            out = tuple(x for x in out if x is not _not_made)
+            out = Union[out]
         else:
-            raise ValueError(f"array_name_format {_array_name_format} not recognised")
-        out = _MetaAbstractArray(
-            name,
-            (AbstractArray,),
-            dict(
-                array_type=array_type,
-                dtypes=cls.dtypes,
-                dims=dims,
-                index_variadic=index_variadic,
-            ),
-        )
-        if getattr(typing, "GENERATING_DOCUMENTATION", False):
-            out.__module__ = "builtins"
-        else:
-            out.__module__ = "jaxtyping"
+            out = _make(array_type)
         return out
 
 
@@ -425,7 +541,9 @@ class AbstractDtype(metaclass=_MetaAbstractDtype):
 
         dtypes: Union[Literal[_any_dtype], str, List[str]] = cls.dtypes
         if isinstance(dtypes, str):
-            dtypes = [dtypes]
+            dtypes = (dtypes,)
+        elif dtypes is not _any_dtype:
+            dtypes = tuple(dtypes)
         cls.dtypes = dtypes
 
 
@@ -459,7 +577,8 @@ if TYPE_CHECKING:
         Annotated as UInt64,
     )
 else:
-    _bool = "bool_"
+    _bool = "bool"
+    _bool_ = "bool_"
     _uint8 = "uint8"
     _uint16 = "uint16"
     _uint32 = "uint32"
@@ -502,6 +621,7 @@ else:
     Complex64 = _make_dtype(_complex64, "Complex64")
     Complex128 = _make_dtype(_complex128, "Complex128")
 
+    bools = [_bool, _bool_]
     uints = [_uint8, _uint16, _uint32, _uint64]
     ints = [_int8, _int16, _int32, _int64]
     floats = [_bfloat16, _float16, _float32, _float64]
@@ -510,7 +630,7 @@ else:
     # We match NumPy's type hierarachy in what types to provide. See the diagram at
     # https://numpy.org/doc/stable/reference/arrays.scalars.html#scalars
 
-    Bool = _make_dtype(_bool, "Bool")
+    Bool = _make_dtype(bools, "Bool")
     UInt = _make_dtype(uints, "UInt")
     Int = _make_dtype(ints, "Int")
     Integer = _make_dtype(uints + ints, "Integer")

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()

setup.py

@@ -63,7 +63,7 @@ classifiers = [
     "Topic :: Scientific/Engineering :: Mathematics",
 ]
 
-python_requires = "~=3.7"
+python_requires = "~=3.8"
 
 # We use typeguard internally (in a fairly minimal way), but it's not required that
 # end users make the same choice.

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_array.py

@@ -17,11 +17,14 @@
 # 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.
 
+from typing import get_args, get_origin, Union
+
 import jax.numpy as jnp
 import jax.random as jr
+import numpy as np
 import pytest
 
-from jaxtyping import AbstractDtype, Array, Float, Float32, jaxtyped, Shaped
+from jaxtyping import AbstractDtype, Array, ArrayLike, Float, Float32, jaxtyped, Shaped
 
 from .helpers import ParamError, ReturnError
 
@@ -409,3 +412,54 @@ def test_incomplete_symbolic(typecheck, getkey):
     x = jr.normal(getkey(), (4,))
     with pytest.raises(NameError):
         foo(x)
+
+
+def test_arraylike(typecheck, getkey):
+    floatlike1 = Float32[ArrayLike, ""]
+    floatlike2 = Float[ArrayLike, ""]
+    floatlike3 = Float32[ArrayLike, "4"]
+
+    assert get_origin(floatlike1) is Union
+    assert get_origin(floatlike2) is Union
+    assert get_origin(floatlike3) is Union
+    assert set(get_args(floatlike1)) == {
+        Float32[Array, ""],
+        Float32[np.ndarray, ""],
+        Float32[np.bool_, ""],
+        Float32[np.number, ""],
+        float,
+    }
+    assert set(get_args(floatlike2)) == {
+        Float[Array, ""],
+        Float[np.ndarray, ""],
+        Float[np.bool_, ""],
+        Float[np.number, ""],
+        float,
+    }
+    assert set(get_args(floatlike3)) == {
+        Float32[Array, "4"],
+        Float32[np.ndarray, "4"],
+        Float32[np.bool_, "4"],
+        Float32[np.number, "4"],
+    }
+
+    shaped1 = Shaped[ArrayLike, ""]
+    shaped2 = Shaped[ArrayLike, "4"]
+    assert get_origin(shaped1) is Union
+    assert get_origin(shaped2) is Union
+    assert set(get_args(shaped1)) == {
+        Shaped[Array, ""],
+        Shaped[np.ndarray, ""],
+        Shaped[np.bool_, ""],
+        Shaped[np.number, ""],
+        bool,
+        int,
+        float,
+        complex,
+    }
+    assert set(get_args(shaped2)) == {
+        Shaped[Array, "4"],
+        Shaped[np.ndarray, "4"],
+        Shaped[np.bool_, "4"],
+        Shaped[np.number, "4"],
+    }

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__