diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py index 4cceec7..2700aed 100644 --- a/jaxtyping/array_types.py +++ b/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): diff --git a/setup.py b/setup.py index 1a97619..37b29b3 100644 --- a/setup.py +++ b/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"]) diff --git a/test/test_array.py b/test/test_array.py index 84130ec..eb7d57c 100644 --- a/test/test_array.py +++ b/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):