diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py index 8a9f066..2d69fc0 100644 --- a/jaxtyping/array_types.py +++ b/jaxtyping/array_types.py @@ -253,7 +253,7 @@ def _make_metaclass(base_metaclass): def _check_scalar(dtype, dtypes, dims): if len(dims) != 0: - return False + return dims == (_anonymous_variadic_dim,) return (_any_dtype is dtypes) or any(d.startswith(dtype) for d in dtypes) @@ -469,6 +469,8 @@ class _MetaAbstractDtype(type): out = Union[out] else: out = _make_array(array_type, dim_str, cls.dtypes, cls.__name__) + if out is _not_made: + raise ValueError("Invalid jaxtyping type annotation.") return out diff --git a/jaxtyping/decorator.py b/jaxtyping/decorator.py index bb1ac6d..e260480 100644 --- a/jaxtyping/decorator.py +++ b/jaxtyping/decorator.py @@ -77,8 +77,10 @@ def _jaxtyped_typechecker(typechecker): def _wrapper(kls): assert inspect.isclass(kls) if dataclasses.is_dataclass(kls): - init = jaxtyped(typechecker(kls.__init__)) - kls.__init__ = init + if type(kls.__init__) is not _Jaxtyped: + # Extra `if` check to work around beartype bug #211 + init = jaxtyped(typechecker(kls.__init__)) + kls.__init__ = init return kls return _wrapper