Tidied up _#* into acting as prefixes, regardless of order. Closes #10

.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

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
 

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"

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}']"

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,))