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, "..."]