Fixed Float[ArrayLike, "#*foo"] leaving out bool/int/float/complex.

jaxtyping/_array_types.py

@@ -291,8 +291,11 @@ def _make_metaclass(base_metaclass):
 
 
 def _check_scalar(dtype, dtypes, dims):
-    if len(dims) != 0:
-        return dims == (_anonymous_variadic_dim,)
+    for dim in dims:
+        if dim is not _anonymous_variadic_dim and not isinstance(
+            dim, _NamedVariadicDim
+        ):
+            return False
     return (_any_dtype is dtypes) or any(d.startswith(dtype) for d in dtypes)
 
 

test/test_array.py

@@ -575,3 +575,17 @@ def test_extension(typecheck, getkey):
     g(jr.split(jr.PRNGKey(0)))
     with pytest.raises(ParamError):
         g(jr.split(jr.PRNGKey(0), 3))
+
+
+def test_scalar_variadic_dim():
+    assert Float[float, "..."] is float
+    assert Float[float, "#*shape"] is float
+
+    # This one is a bit weird -- it should really also assert that shape==(), but we
+    # don't implement that.
+    assert Float[float, "*shape"] is float
+
+
+def test_scalar_dtype_mismatch():
+    with pytest.raises(ValueError):
+        Float[bool, "..."]