From 1a048b1f2fa757b9e2db1fb70cc4fcccca1fe86d Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Sat, 21 Oct 2023 11:49:28 -0700
Subject: [PATCH] Fixed Float[ArrayLike, "#*foo"] leaving out
 bool/int/float/complex.

---
 jaxtyping/_array_types.py |  7 +++++--
 test/test_array.py        | 14 ++++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/jaxtyping/_array_types.py b/jaxtyping/_array_types.py
index babdb76..06708fc 100644
--- a/jaxtyping/_array_types.py
+++ b/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)
 
 
diff --git a/test/test_array.py b/test/test_array.py
index 28fc966..c52a53e 100644
--- a/test/test_array.py
+++ b/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, "..."]