Broadcast+variadic is now more lenient, and supports anything broadcastable together. Closes #8

jaxtyping/array_types.py

@@ -22,6 +22,7 @@ from typing import Any, Dict, List, NoReturn, Optional, Tuple, Union
 from typing_extensions import Literal
 
 import jax.numpy as jnp
+import numpy as np
 
 from .decorator import storage
 
@@ -142,29 +143,34 @@ class _MetaAbstractArray(type):
             if j is not None and not _check_dims(cls.dims[j:], obj.shape[j:], memo):
                 return False
             variadic_dim = cls.dims[i]
-            if variadic_dim is not _anonymous_variadic_dim:
+            if variadic_dim is _anonymous_variadic_dim:
+                return True
+            else:
                 variadic_name = variadic_dim.name
                 try:
-                    variadic_shape = memo[variadic_name]
+                    if variadic_dim.broadcastable:
+                        variadic_shapes = memo[variadic_name]
+                    else:
+                        variadic_shape = memo[variadic_name]
                 except KeyError:
-                    memo[variadic_name] = obj.shape[i:j]
+                    if variadic_dim.broadcastable:
+                        memo[variadic_name] = [obj.shape[i:j]]
+                    else:
+                        memo[variadic_name] = obj.shape[i:j]
+                    return True
                 else:
                     if variadic_dim.broadcastable:
-                        new_variadic_shape = []
-                        obj_shape = obj.shape[i:j]
-                        if len(variadic_shape) != len(obj_shape):
-                            return False
-                        for old_size, new_size in zip(variadic_shape, obj_shape):
-                            if old_size == 1:
-                                new_variadic_shape.append(new_size)
-                            else:
-                                if new_size != 1 and old_size != new_size:
-                                    return False
-                                new_variadic_shape.append(old_size)
-                        memo[variadic_name] = tuple(new_variadic_shape)
+                        new_shape = obj.shape[i:j]
+                        for existing_shape in variadic_shapes:
+                            try:
+                                np.broadcast_shapes(new_shape, existing_shape)
+                            except ValueError:
+                                return False
+                        variadic_shapes.append(new_shape)
+                        return True
                     else:
                         return variadic_shape == obj.shape[i:j]
-            return True
+            assert False
 
 
 class AbstractArray(metaclass=_MetaAbstractArray):

setup.py

@@ -28,7 +28,8 @@ _here = pathlib.Path(__file__).resolve().parent
 
 name = "jaxtyping"
 
-# for simplicity we actually store the version in the __version__ attribute in the source
+# for simplicity we actually store the version in the __version__ attribute in the
+# source
 with open(_here / name / "__init__.py") as f:
     meta_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]", f.read(), re.M)
     if meta_match:
@@ -40,7 +41,10 @@ author = "Patrick Kidger"
 
 author_email = "contact@kidger.site"
 
-description = "Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees."
+description = (
+    "Type annotations and runtime checking for shape and dtype of JAX "
+    "arrays, and PyTrees."
+)
 
 with open(_here / "README.md", "r") as f:
     readme = f.read()
@@ -63,7 +67,12 @@ python_requires = "~=3.7"
 
 # We use typeguard internally (in a fairly minimal way), but it's not required that
 # end users make the same choice.
-install_requires = ["jax>=0.3.4", "typeguard>=2.13.3", "typing_extensions>=4.2.0"]
+install_requires = [
+    "jax>=0.3.4",
+    "numpy>=1.20.0",
+    "typeguard>=2.13.3",
+    "typing_extensions>=4.2.0",
+]
 
 entry_points = dict(pytest11=["jaxtyping = jaxtyping.pytest_plugin"])
 

test/test_array.py

@@ -283,12 +283,11 @@ def test_broadcast_variadic_named(typecheck, getkey):
     g(b, b)
     g(c, c)
     g(d, d)
+    g(b, c)
     with pytest.raises(ParamError):
         g(a, b)
     with pytest.raises(ParamError):
         g(a, c)
-    with pytest.raises(ParamError):
-        g(b, c)
     with pytest.raises(ParamError):
         g(a, b)
     with pytest.raises(ParamError):
@@ -296,26 +295,20 @@ def test_broadcast_variadic_named(typecheck, getkey):
 
     g(a, j)
     g(b, j)
-    with pytest.raises(ParamError):
-        g(c, j)
-    with pytest.raises(ParamError):
-        g(d, j)
-    with pytest.raises(ParamError):
-        g(b, k)
+    g(c, j)
+    g(d, j)
+    g(b, k)
     g(c, k)
     with pytest.raises(ParamError):
         g(d, k)
     with pytest.raises(ParamError):
         g(c, l)
     g(d, l)
-    with pytest.raises(ParamError):
-        g(a, m)
+    g(a, m)
     g(c, m)
     g(d, m)
-    with pytest.raises(ParamError):
-        g(a, n)
-    with pytest.raises(ParamError):
-        g(b, n)
+    g(a, n)
+    g(b, n)
     with pytest.raises(ParamError):
         g(c, n)
     with pytest.raises(ParamError):