diff --git a/.flake8 b/.flake8 index 7b18abd..8026b92 100644 --- a/.flake8 +++ b/.flake8 @@ -1,4 +1,4 @@ [flake8] -max-line-length = 120 +max-line-length = 88 ignore = W291,W293,W503,W504,E123,E126,E203,E402,E701,E731,F722 per-file-ignores = __init__.py: F401 diff --git a/API.md b/API.md index bcf8ce1..aae98bf 100644 --- a/API.md +++ b/API.md @@ -6,22 +6,27 @@ Each array is denoted by a type `dtype[shape]`, such as `f32["batch channels"]`. ### Shape -The shape should be a string of space-separated symbols, such as "a b c d". Each symbol can be: +The shape should be a string of space-separated symbols, such as "a b c d". Each symbol can be either an: - `int`: fixed-size axis, e.g. `f32["28 28"]`. - `str`: variable-size axis, e.g. `f32["channels"]`. -- `_`: anonymous axis, e.g. `f32["batch channels _ _"]`. -- `...`: anonymous zero or more axes, e.g. `f32["... c h w"]` -- `*name`: zero or more variable-size axes, e.g. `f32["*batch c h w"]` -- Append `#` to a dimension size to indicate that it can be that size *or* equal to one -- i.e. broadcasting is acceptable. When calling a function, variable-size axes will be matched up across all arguments and checked for consistency. (See [runtime type checking](#runtime-type-checking) below.) +In addition some modifiers can be applied: +- Prepend `*` to a dimension to indicate that it can match multiple axes, e.g. `f32["*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. `add(x: f32["#foo"], y: f32["#foo"]) -> f32["#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. `f32["b c _ _"]`. + +The order of these modifiers does not matter. + +As a special case: +- `...`: anonymous zero or more axes (equivalent to `*_`) e.g. `f32["... c h w"]` + Some notes: - To denote a scalar shape use `""`, e.g. `f32[""]`. - To denote an arbitrary shape (and only check dtype) use `"..."`, e.g. `f32["..."]`. -- You cannot have multiple variadic axes, i.e. you can only use `...` or `*name` at most once in each array. -- An example of broadcasting in one dimension: `add(x: f32["foo#"], y: f32["foo#"]) -> f32["foo#"]`. -- An example of broadcasting multiple dimensions: `add(x: f32["*foo#"], y: f32["*foo#"]) -> f32["*foo#"]`. +- You cannot have more than one use of multiple-axes, i.e. you can only use `...` or `*name` at most once in each array. +- An example of broadcasting multiple dimensions: `add(x: f32["*#foo"], y: f32["*#foo"]) -> f32["*#foo"]`. ### Dtype diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index bf4a5ad..ad26ecd 100644 --- a/jaxtyping/__init__.py +++ b/jaxtyping/__init__.py @@ -51,4 +51,4 @@ from .import_hook import install_import_hook from .pytree_type import PyTree -__version__ = "0.0.2" +__version__ = "0.0.3" diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py index 348ee06..4cceec7 100644 --- a/jaxtyping/array_types.py +++ b/jaxtyping/array_types.py @@ -185,7 +185,8 @@ class _MetaAbstractDtype(type): def __getitem__(cls, dim_str: str) -> _MetaAbstractArray: if not isinstance(dim_str, str): raise ValueError( - "Shape specification must be a string. Axes should be separated with spaces." + "Shape specification must be a string. Axes should be separated with " + "spaces." ) dims = [] index_variadic = None @@ -195,29 +196,100 @@ class _MetaAbstractDtype(type): raise ValueError( "Dimensions should be separated with spaces, not commas" ) - broadcastable = False if elem.endswith("#"): - broadcastable = True - elem = elem[:-1] - try: - elem = int(elem) - except ValueError: - if elem == "_": - elem = _anonymous_dim - elif elem == "...": - if index_variadic is not None: - raise ValueError("Cannot have multiple variadic dimensions") - index_variadic = index - elem = _anonymous_variadic_dim - elif elem[0] == "*": - if index_variadic is not None: - raise ValueError("Cannot have multiple variadic dimensions") - index_variadic = index - elem = _NamedVariadicDim(elem[1:], broadcastable) - else: - elem = _NamedDim(elem, broadcastable) + raise ValueError( + "As of jaxtyping v0.0.3, broadcastable dimensions are now denoted " + "with a # at the start, rather than at the end" + ) + + if "..." in elem: + if elem != "...": + raise ValueError( + "Anonymous multiple dimension '...' must be used on its own; " + f"got {elem}" + ) + broadcastable = False + variadic = True + anonymous = True + is_fixed = False else: + broadcastable = False + variadic = False + anonymous = False + while True: + if len(elem) == 0: + # This branch needed as just `_` is valid + break + first_char = elem[0] + if first_char == "#": + if broadcastable: + raise ValueError( + "Do not use # twice to denote broadcastability, e.g. " + "`##foo` is not allowed" + ) + broadcastable = True + elem = elem[1:] + elif first_char == "*": + if variadic: + raise ValueError( + "Do not use * twice to denote accepting multiple " + "dimensions, e.g. `**foo` is not allowed" + ) + variadic = True + elem = elem[1:] + elif first_char == "_": + if anonymous: + raise ValueError( + "Do not use _ twice to denote anonymity, e.g. `__foo` " + "is not allowed" + ) + anonymous = True + elem = elem[1:] + else: + break + try: + elem = int(elem) + except ValueError: + is_fixed = False + else: + is_fixed = True + + if variadic: + if index_variadic is not None: + raise ValueError( + "Cannot use multiple-dimension specifiers (`*name` or `...`) " + "more than once" + ) + index_variadic = index + + if is_fixed: + if variadic: + raise ValueError( + "Cannot have a fixed axis bind to multiple dimensions, e.g. " + "`*4` is not allowed" + ) + if anonymous: + raise ValueError( + "Cannot have a fixed axis be anonymous, e.g. `_4` is not " + "allowed" + ) elem = _FixedDim(elem, broadcastable) + else: + if anonymous: + if broadcastable: + raise ValueError( + "Cannot have a dimension be both anonymous and " + "broadcastable, e.g. `#_` is not allowed" + ) + if variadic: + elem = _anonymous_variadic_dim + else: + elem = _anonymous_dim + else: + if variadic: + elem = _NamedVariadicDim(elem, broadcastable) + else: + elem = _NamedDim(elem, broadcastable) dims.append(elem) if _array_name_format == "dtype_and_shape": name = f"{cls.__name__}['{dim_str}']" diff --git a/test/test_array.py b/test/test_array.py index e6de38b..84130ec 100644 --- a/test/test_array.py +++ b/test/test_array.py @@ -227,7 +227,7 @@ def test_anonymous_variadic(typecheck, getkey): def test_broadcast_fixed(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32["4#"]): + def g(x: f32["#4"]): pass g(jr.normal(getkey(), (4,))) @@ -240,7 +240,7 @@ def test_broadcast_fixed(typecheck, getkey): def test_broadcast_named(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32[" foo#"], y: f32[" foo#"]): + def g(x: f32[" #foo"], y: f32[" #foo"]): pass a = jr.normal(getkey(), (3,)) @@ -264,7 +264,7 @@ def test_broadcast_named(typecheck, getkey): def test_broadcast_variadic_named(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32[" *foo#"], y: f32[" *foo#"]): + def g(x: f32[" *#foo"], y: f32[" *#foo"]): pass a = jr.normal(getkey(), (3,))