From 38be24f9c8e7cedc8d76980bc832ecdae21c34cb Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Sun, 5 Mar 2023 20:09:49 -0800
Subject: [PATCH] beartype+inheritance fix. Bool[int, '...'] now correctly
 raises an error.

---
 jaxtyping/array_types.py | 4 +++-
 jaxtyping/decorator.py   | 6 ++++--
 2 files changed, 7 insertions(+), 3 deletions(-)

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