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.

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