mirror of
https://github.com/wassname/jaxtyping.git
synced 2026-09-12 12:32:19 +08:00
beartype+inheritance fix. Bool[int, '...'] now correctly raises an error.
This commit is contained in:
@@ -253,7 +253,7 @@ def _make_metaclass(base_metaclass):
|
|||||||
|
|
||||||
def _check_scalar(dtype, dtypes, dims):
|
def _check_scalar(dtype, dtypes, dims):
|
||||||
if len(dims) != 0:
|
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)
|
return (_any_dtype is dtypes) or any(d.startswith(dtype) for d in dtypes)
|
||||||
|
|
||||||
|
|
||||||
@@ -469,6 +469,8 @@ class _MetaAbstractDtype(type):
|
|||||||
out = Union[out]
|
out = Union[out]
|
||||||
else:
|
else:
|
||||||
out = _make_array(array_type, dim_str, cls.dtypes, cls.__name__)
|
out = _make_array(array_type, dim_str, cls.dtypes, cls.__name__)
|
||||||
|
if out is _not_made:
|
||||||
|
raise ValueError("Invalid jaxtyping type annotation.")
|
||||||
return out
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -77,8 +77,10 @@ def _jaxtyped_typechecker(typechecker):
|
|||||||
def _wrapper(kls):
|
def _wrapper(kls):
|
||||||
assert inspect.isclass(kls)
|
assert inspect.isclass(kls)
|
||||||
if dataclasses.is_dataclass(kls):
|
if dataclasses.is_dataclass(kls):
|
||||||
init = jaxtyped(typechecker(kls.__init__))
|
if type(kls.__init__) is not _Jaxtyped:
|
||||||
kls.__init__ = init
|
# Extra `if` check to work around beartype bug #211
|
||||||
|
init = jaxtyped(typechecker(kls.__init__))
|
||||||
|
kls.__init__ = init
|
||||||
return kls
|
return kls
|
||||||
|
|
||||||
return _wrapper
|
return _wrapper
|
||||||
|
|||||||
Reference in New Issue
Block a user