Compare commits

..
2 Commits
Author SHA1 Message Date
Patrick Kidger 165065756f Static type-checking fix (#34) 2022-09-24 18:27:22 -07:00
ebrevdo dcd73e3431 Add support for e.g. jaxtyping.Float[Union[...], ...] in py3.8 (#31)
* Add support for e.g. jaxtyping.Float[Union[...], ...] in py3.8

Turns out that python3.8, Union lacks the __name__ attribute.  Use repr()
in these cases.

* Fix linter.

* Remove implicit cast to bool in favor of try/except.
2022-09-22 13:25:29 -07:00
2 changed files with 18 additions and 9 deletions
+12 -8
View File
@@ -20,14 +20,18 @@
import typing
if getattr(typing, "GENERATING_DOCUMENTATION", False):
class Array:
pass
Array.__module__ = "builtins"
else:
if typing.TYPE_CHECKING:
# type checkers don't know which branch below will be executed
from jax.numpy import ndarray as Array
else:
if getattr(typing, "GENERATING_DOCUMENTATION", False):
class Array:
pass
Array.__module__ = "builtins"
else:
from jax.numpy import ndarray as Array
from .array_types import (
AbstractArray,
@@ -63,4 +67,4 @@ from .import_hook import install_import_hook
from .pytree_type import PyTree
__version__ = "0.2.5"
__version__ = "0.2.6"
+6 -1
View File
@@ -384,7 +384,12 @@ class _MetaAbstractDtype(type):
elem = _SymbolicDim(elem, broadcastable)
dims.append(elem)
if _array_name_format == "dtype_and_shape":
name = f"{cls.__name__}[{array_type.__name__}, '{dim_str}']"
# In python 3.8, e.g., typing.Union lacks `__name__`.
try:
type_str = array_type.__name__
except AttributeError:
type_str = repr(array_type)
name = f"{cls.__name__}[{type_str}, '{dim_str}']"
elif _array_name_format == "array":
name = array_type.__name__
else: