diff --git a/API.md b/API.md index 4e15298..9b390a9 100644 --- a/API.md +++ b/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. diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index 87c285d..5f2bc36 100644 --- a/jaxtyping/__init__.py +++ b/jaxtyping/__init__.py @@ -114,4 +114,4 @@ elif has_jax: del has_jax -__version__ = "0.2.14" +__version__ = "0.2.15" diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py index 2d69fc0..4be066b 100644 --- a/jaxtyping/array_types.py +++ b/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: - # Common mistake + if "," in elem and "(" not in elem: + # 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: - elem = int(elem) - except ValueError: - if len(elem) == 0 or elem.isidentifier(): - dim_type = _DimType.named - else: - dim_type = _DimType.symbolic + if len(elem) == 0 or elem.isidentifier(): + dim_type = _DimType.named 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) diff --git a/test/test_array.py b/test/test_array.py index 9653e5b..0463d05 100644 --- a/test/test_array.py +++ b/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))