diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index e66d044..d6ecfca 100644 --- a/jaxtyping/__init__.py +++ b/jaxtyping/__init__.py @@ -65,6 +65,8 @@ elif has_jax: if typing.TYPE_CHECKING: # Introduce an indirection so that we can `import X as X` to make it clear that # these are public. + from jax.typing import DTypeLike as DTypeLike + from ._indirection import ( BFloat16 as BFloat16, Bool as Bool, @@ -121,8 +123,13 @@ else: ) if has_jax: + import jax.typing + from ._array_types import Key as Key + if hasattr(jax.typing, "DTypeLike"): + from jax.typing import DTypeLike as DTypeLike + # Now import PyTreeDef and PyTree if typing.TYPE_CHECKING: