Some improvements (#77)

* Various improvements.

- Added support for functions in symbolic dimensions, e.g. "min(foo,bar)", which were previously disallowed due to the presence of a comma. (#51)
- Added support for adding ignored names to dimensions, e.g. "cols=4". (#76)

* Now works with Python 3.10 A | B union types.

.flake8

@@ -1,4 +0,0 @@
-[flake8]
-max-line-length = 88
-ignore = W291,W293,W503,W504,E123,E126,E203,E402,E701,E731,F722
-per-file-ignores = __init__.py: F401

.isort.cfg

@@ -1,7 +0,0 @@
-[settings]
-force_alphabetical_sort_within_sections=true
-lines_after_imports=2
-profile=black
-combine_as_imports=True
-treat_comments_as_code=true
-extra_standard_library=typing_extensions

.pre-commit-config.yaml

@@ -22,17 +22,7 @@ repos:
     rev: 22.3.0
     hooks:
     - id: black
-  - repo: https://github.com/nbQA-dev/nbQA
+  - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: 1.6.3
+    rev: 'v0.0.255'
     hooks:
-    - id: nbqa-black
+      - id: ruff
-    - id: nbqa-isort 
-    - id: nbqa-flake8
-  - repo: https://github.com/PyCQA/isort
-    rev: 5.12.0
-    hooks:
-    - id: isort
-  - repo: https://github.com/pycqa/flake8
-    rev: 4.0.1
-    hooks:
-    - id: flake8

API.md

@@ -17,6 +17,7 @@ In addition some modifiers can be applied:
 - Prepend `*` to a dimension to indicate that it can match multiple axes, e.g. `"*batch c h w"` will match zero or more batch axes.
 - Prepend `#` to a dimension to indicate that it can be that size *or* equal to one -- i.e. broadcasting is acceptable, e.g. `def add(x: Float[Array, "#foo"], y: Float[Array, "#foo"]) -> Float[Array, "#foo"]`.
 - Prepend `_` to a dimension to disable any runtime checking of that dimension (so that it can be used just as documentation). This can also be used as just `_` on its own: e.g. `"b c _ _"`.
+- Documentation-only names (i.e. they're ignored by jaxtyping) can be handled by prepending a name followed by `=` e.g. `Float[Array, "rows=4 cols=3"]`.
 
 When using multiple modifiers, their order does not matter.
 

FAQ.md

@@ -12,13 +12,13 @@ jaxtyping and `jax.jit` synergise beautifully.
 
 When calling JAX operations wrapped in a `jax.jit`, then the dtype/shape-checking will happen at trace time. (When JAX traces your function prior to compiling it.) The actual compiled code does not have any dtype/shape-checking, and will therefore still be just as fast as before!
 
-## `flake8` is throwing an error.
+## `flake8` or Ruff are throwing an error.
 
 In type annotations, strings are used for two different things. Sometimes they're strings. Sometimes they're "forward references", used to refer to a type that will be defined later.
 
 Some tooling in the Python ecosystem assumes that only the latter is true, and will throw spurious errors if you try to use a string just as a string (like we do).
 
-In the case of `flake8`, at least, this is easily resolved. Multi-dimensional arrays (e.g. `Float32[Array, "b c"]`) will throw a very unusual error (F722, syntax error in forward annotation), so you can safely just disable this particular error globally. Uni-dimensional arrays (e.g. `Float32[Array, "x"]`) will throw an error that's actually useful (F821, undefined name), so instead of disabling this globally, you should instead prepend a space to the start of your shape, e.g. `Float32[Array, " x"]`. `jaxtyping` will treat this in the same way, whilst `flake8` will now throw an F722 error that you can disable as before.
+In the case of `flake8`, or Ruff, this can be resolved. Multi-dimensional arrays (e.g. `Float32[Array, "b c"]`) will throw a very unusual error (F722, syntax error in forward annotation), so you can safely just disable this particular error globally. Uni-dimensional arrays (e.g. `Float32[Array, "x"]`) will throw an error that's actually useful (F821, undefined name), so instead of disabling this globally, you should instead prepend a space to the start of your shape, e.g. `Float32[Array, " x"]`. `jaxtyping` will treat this in the same way, whilst `flake8` will now throw an F722 error that you can disable as before.
 
 ## Does jaxtyping use [PEP 646](https://www.python.org/dev/peps/pep-0646/) (variadic generics)?
 

jaxtyping/__init__.py

@@ -110,8 +110,8 @@ if typing.TYPE_CHECKING:
     # 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
+    from .pytree_type import PyTree as PyTree  # noqa: F401
 
 del has_jax
 
-__version__ = "0.2.14"
+__version__ = "0.2.15"

jaxtyping/array_types.py

@@ -19,6 +19,8 @@
 
 import enum
 import functools as ft
+import sys
+import types
 import typing
 from typing import (
     Any,
@@ -267,6 +269,11 @@ class AbstractArray(metaclass=_MetaAbstractArray):
 _not_made = object()
 
 
+_union_types = [typing.Union]
+if sys.version_info >= (3, 10):
+    _union_types.append(types.UnionType)
+
+
 @ft.lru_cache(maxsize=None)
 def _make_array(array_type, dim_str, dtypes, name):
     if not isinstance(dim_str, str):
@@ -277,8 +284,10 @@ def _make_array(array_type, dim_str, dtypes, name):
     dims = []
     index_variadic = None
     for index, elem in enumerate(dim_str.split()):
-        if "," in elem:
+        if "," in elem and "(" not in elem:
-            # Common mistake
+            # Common mistake.
+            # Disable in the case that there's brackets to allow for function calls,
+            # e.g. `min(foo,bar)`, in symbolic dimensions.
             raise ValueError("Dimensions should be separated with spaces, not commas")
         if elem.endswith("#"):
             raise ValueError(
@@ -329,17 +338,21 @@ def _make_array(array_type, dim_str, dtypes, name):
                         )
                     anonymous = True
                     elem = elem[1:]
+                # Allow e.g. `foo=4` as an alternate syntax for just `4`, so that one
+                # can write e.g. `Float[Array, "rows=3 cols=4"]`
+                elif elem.count("=") == 1:
+                    _, elem = elem.split("=")
                 else:
                     break
-            try:
+            if len(elem) == 0 or elem.isidentifier():
-                elem = int(elem)
+                dim_type = _DimType.named
-            except ValueError:
-                if len(elem) == 0 or elem.isidentifier():
-                    dim_type = _DimType.named
-                else:
-                    dim_type = _DimType.symbolic
             else:
-                dim_type = _DimType.fixed
+                try:
+                    elem = int(elem)
+                except ValueError:
+                    dim_type = _DimType.symbolic
+                else:
+                    dim_type = _DimType.fixed
 
         if variadic:
             if index_variadic is not None:
@@ -460,7 +473,7 @@ class _MetaAbstractDtype(type):
             )
         array_type, dim_str = item
         del item
-        if typing.get_origin(array_type) is typing.Union:
+        if typing.get_origin(array_type) in _union_types:
             out = [
                 _make_array(x, dim_str, cls.dtypes, cls.__name__)
                 for x in typing.get_args(array_type)

jaxtyping/decorator.py

@@ -21,40 +21,20 @@ import dataclasses
 import functools as ft
 import inspect
 import threading
+import types
 import weakref
 
 
 storage = threading.local()
 
 
-_fns = weakref.WeakKeyDictionary()
+_jaxtyped_fns = weakref.WeakSet()
-
-
-class _Jaxtyped:
-    def __init__(self, fn):
-        # Stored externally so that it doesn't get blatted in the `ft.wraps` below by
-        # a function that already has a `fn` attribute.
-        _fns[self] = fn
-
-    def __get__(self, instance, owner):
-        fn = _fns[self]
-        return ft.wraps(fn)(_Jaxtyped(fn.__get__(instance, owner)))
-
-    def __call__(self, *args, **kwargs):
-        try:
-            memo_stack = storage.memo_stack
-        except AttributeError:
-            memo_stack = storage.memo_stack = []
-        memo_stack.append(({}, {}, {}))
-        fn = _fns[self]
-        try:
-            return fn(*args, **kwargs)
-        finally:
-            memo_stack.pop()
 
 
 def jaxtyped(fn):
-    if inspect.isclass(fn):  # allow decorators on class definitions
+    if type(fn) is types.FunctionType and fn in _jaxtyped_fns:
+        return fn
+    elif inspect.isclass(fn):  # allow decorators on class definitions
         if dataclasses.is_dataclass(fn):
             init = jaxtyped(fn.__init__)
             fn.__init__ = init
@@ -63,8 +43,44 @@ def jaxtyped(fn):
             raise ValueError(
                 "jaxtyped may only be added as a class decorator to dataclasses"
             )
+    # It'd be lovely if we could handle arbitrary descriptors, and not just the builtin
+    # ones. Unfortunately that means returning a class instance with a __get__ method,
+    # and that turns out to break loads of other things. See beartype issue #211 and
+    # jaxtyping issue #71.
+    elif isinstance(fn, classmethod):
+        return classmethod(jaxtyped(fn.__func__))
+    elif isinstance(fn, staticmethod):
+        return staticmethod(jaxtyped(fn.__func__))
+    elif isinstance(fn, property):
+        if fn.fget is None:
+            fget = None
+        else:
+            fget = jaxtyped(fn.fget)
+        if fn.fset is None:
+            fset = None
+        else:
+            fset = jaxtyped(fn.fset)
+        if fn.fdel is None:
+            fdel = None
+        else:
+            fdel = jaxtyped(fn.fdel)
+        return property(fget=fget, fset=fset, fdel=fdel)
     else:
-        return ft.wraps(fn)(_Jaxtyped(fn))
+
+        @ft.wraps(fn)
+        def wrapped_fn(*args, **kwargs):
+            try:
+                memo_stack = storage.memo_stack
+            except AttributeError:
+                memo_stack = storage.memo_stack = []
+            memo_stack.append(({}, {}, {}))
+            try:
+                return fn(*args, **kwargs)
+            finally:
+                memo_stack.pop()
+
+        _jaxtyped_fns.add(wrapped_fn)
+        return wrapped_fn
 
 
 def _jaxtyped_typechecker(typechecker):
@@ -77,10 +93,8 @@ def _jaxtyped_typechecker(typechecker):
     def _wrapper(kls):
         assert inspect.isclass(kls)
         if dataclasses.is_dataclass(kls):
-            if type(kls.__init__) is not _Jaxtyped:
+            init = jaxtyped(typechecker(kls.__init__))
-                # Extra `if` check to work around beartype bug #211
+            kls.__init__ = init
-                init = jaxtyped(typechecker(kls.__init__))
-                kls.__init__ = init
         return kls
 
     return _wrapper

pyproject.toml

@@ -0,0 +1,10 @@
+[tool.ruff]
+select = ["E", "F", "I001"]
+ignore = ["E721", "E731", "F722"]
+ignore-init-module-imports = true
+
+[tool.ruff.isort]
+combine-as-imports = true
+lines-after-imports = 2
+extra-standard-library = ["typing_extensions"]
+order-by-type = false

test/import_hook_tester_transitive/__init__.py

@@ -17,4 +17,4 @@
 # 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 . import another_file
+from . import another_file  # noqa: F401

test/import_hook_tester_typeguard.py

@@ -17,9 +17,7 @@
 # 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
 
@@ -37,26 +35,29 @@ with pytest.raises(ParamError):
     g(jnp.array(1))
 
 
-class M(eqx.Module):
+# Typeguard 3.0 no longer supports this.
-    foo: int
+#
-    bar: Float32[jnp.ndarray, " a"]
+# class M(eqx.Module):
-
+#     foo: int
-
+#     bar: Float32[jnp.ndarray, " a"]
-M(1, jnp.array([1.0]))
+#
-with pytest.raises(ParamError):
+#
-    M(1.0, jnp.array([1.0]))
+# M(1, jnp.array([1.0]))
-with pytest.raises(ParamError):
+# with pytest.raises(ParamError):
-    M(1, jnp.array(1.0))
+#     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"]
+# @dataclasses.dataclass
-
+# class D:
-
+#     foo: int
-D(1, jnp.array([1.0]))
+#     bar: Float32[jnp.ndarray, " a"]
-with pytest.raises(ParamError):
+#
-    D(1.0, jnp.array([1.0]))
+#
-with pytest.raises(ParamError):
+# D(1, jnp.array([1.0]))
-    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_old.py

@@ -17,9 +17,7 @@
 # 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
 
@@ -37,26 +35,29 @@ with pytest.raises(ParamError):
     g(jnp.array(1))
 
 
-class M(eqx.Module):
+# Typeguard 3.0 no longer supports this.
-    foo: int
+#
-    bar: Float32[jnp.ndarray, " a"]
+# class M(eqx.Module):
-
+#     foo: int
-
+#     bar: Float32[jnp.ndarray, " a"]
-M(1, jnp.array([1.0]))
+#
-with pytest.raises(ParamError):
+#
-    M(1.0, jnp.array([1.0]))
+# M(1, jnp.array([1.0]))
-with pytest.raises(ParamError):
+# with pytest.raises(ParamError):
-    M(1, jnp.array(1.0))
+#     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"]
+# @dataclasses.dataclass
-
+# class D:
-
+#     foo: int
-D(1, jnp.array([1.0]))
+#     bar: Float32[jnp.ndarray, " a"]
-with pytest.raises(ParamError):
+#
-    D(1.0, jnp.array([1.0]))
+#
-with pytest.raises(ParamError):
+# D(1, jnp.array([1.0]))
-    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,6 +17,7 @@
 # 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 sys
 from typing import get_args, get_origin, Union
 
 import jax.numpy as jnp
@@ -470,3 +471,44 @@ def test_subclass():
     assert issubclass(Float[Array, ""], Array)
     assert issubclass(Float[np.ndarray, ""], np.ndarray)
     assert issubclass(Float[torch.Tensor, ""], torch.Tensor)
+
+
+def test_ignored_names():
+    x = Float[np.ndarray, "foo=4"]
+
+    assert isinstance(np.zeros(4), x)
+    assert not isinstance(np.zeros(5), x)
+    assert not isinstance(np.zeros((4, 5)), x)
+
+    y = Float[np.ndarray, "bar qux foo=bar+qux"]
+
+    assert isinstance(np.zeros((2, 3, 5)), y)
+    assert not isinstance(np.zeros((2, 3, 6)), y)
+
+    z = Float[np.ndarray, "bar #foo=bar"]
+
+    assert isinstance(np.zeros((3, 3)), z)
+    assert isinstance(np.zeros((3, 1)), z)
+    assert not isinstance(np.zeros((3, 4)), z)
+
+    # Weird but legal
+    w = Float[np.ndarray, "bar foo=#bar"]
+
+    assert isinstance(np.zeros((3, 3)), w)
+    assert isinstance(np.zeros((3, 1)), w)
+    assert not isinstance(np.zeros((3, 4)), w)
+
+
+def test_symbolic_functions():
+    x = Float[np.ndarray, "foo bar min(foo,bar)"]
+
+    assert isinstance(np.zeros((2, 3, 2)), x)
+    assert isinstance(np.zeros((3, 2, 2)), x)
+    assert not isinstance(np.zeros((3, 2, 4)), x)
+
+
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="requires Python 3.10")
+def test_py310_unions():
+    x = np.zeros(3)
+    y = Shaped[Array | np.ndarray, "_"]
+    assert isinstance(x, get_args(y))

test/test_decorator.py

@@ -4,23 +4,74 @@ from jaxtyping import jaxtyped
 
 
 class M(metaclass=abc.ABCMeta):
+    @jaxtyped
+    def f(self):
+        ...
+
     @jaxtyped
     @classmethod
-    def f(cls):
+    def g1(cls):
         return 3
 
+    @classmethod
+    @jaxtyped
+    def g2(cls):
+        return 4
+
+    @jaxtyped
+    @staticmethod
+    def h1():
+        return 3
+
+    @staticmethod
+    @jaxtyped
+    def h2():
+        return 4
+
     @jaxtyped
     @abc.abstractmethod
-    def g(self):
+    def i1(self):
+        ...
+
+    @abc.abstractmethod
+    @jaxtyped
+    def i2(self):
         ...
 
 
-# Check that the @jaxtyped decorator doesn't blat the __get__ of @classmethod
+class N:
+    @jaxtyped
+    @property
+    def j1(self):
+        return 3
+
+    @property
+    @jaxtyped
+    def j2(self):
+        return 4
+
+
+def test_identity():
+    assert M.f is M.f
+
+
 def test_classmethod():
-    assert M.f() == 3
+    assert M.g1() == 3
+    assert M.g2() == 4
+
+
+def test_staticmethod():
+    assert M.h1() == 3
+    assert M.h2() == 4
 
 
 # Check that the @jaxtyped decorator doesn't blat the __isabstractmethod__ of
 # @abstractmethod
 def test_abstractmethod():
-    assert M.g.__isabstractmethod__
+    assert M.i1.__isabstractmethod__
+    assert M.i2.__isabstractmethod__
+
+
+def test_property():
+    assert N().j1 == 3
+    assert N().j2 == 4

test/test_import_hook.py

@@ -78,7 +78,7 @@ def test_import_hook_beartype_full():
 
 def test_import_hook_transitive():
     hook = install_import_hook(
-        "test.import_hook_tester_transitive", "typeguard.typechecked"
+        "test.import_hook_tester_transitive", "beartype.beartype"
     )
     with hook:
         from . import import_hook_tester_transitive  # noqa: F401