Some improvements (#77)

* Various improvements.

- Added support for functions in symbolic dimensions, e.g. "min(foo,bar)", which were previously disallowed due to the presence of a comma. (#51)
- Added support for adding ignored names to dimensions, e.g. "cols=4". (#76)

* Now works with Python 3.10 A | B union types.

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.
 

jaxtyping/__init__.py

@@ -114,4 +114,4 @@ elif has_jax:
 
 del has_jax
 
-__version__ = "0.2.14"
+__version__ = "0.2.15"

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)

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