diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index 7bce648..e77fb48 100644 --- a/jaxtyping/__init__.py +++ b/jaxtyping/__init__.py @@ -17,16 +17,17 @@ # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. +import functools as ft import importlib.metadata +import importlib.util import typing import warnings +from typing import Union -# First import some things as normal from ._array_types import ( AbstractArray as AbstractArray, AbstractDtype as AbstractDtype, get_array_name_format as get_array_name_format, - has_jax, set_array_name_format as set_array_name_format, ) from ._config import config as config @@ -40,40 +41,68 @@ from ._ipython_extension import load_ipython_extension as load_ipython_extension from ._storage import print_bindings as print_bindings -# Now import Array and ArrayLike if typing.TYPE_CHECKING: - # For imports, we need to explicitly `import X as X` in order for Pyright to see - # them as public. See discussion at https://github.com/microsoft/pyright/issues/2277 + import typing_extensions + from jax import Array as Array - from jax.typing import ArrayLike as ArrayLike -elif has_jax: - if getattr(typing, "GENERATING_DOCUMENTATION", False): + from jax.tree_util import PyTreeDef as PyTreeDef + from jax.typing import ArrayLike as ArrayLike, DTypeLike as DTypeLike - class Array: - pass - - Array.__module__ = "builtins" - - class ArrayLike: - pass - - ArrayLike.__module__ = "builtins" - else: - from jax import Array as Array - - try: - from jax.typing import ArrayLike as ArrayLike - except (ModuleNotFoundError, ImportError): - pass - - -# Import our dtypes -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, + Complex as Complex, + Complex64 as Complex64, + Complex128 as Complex128, + Float as Float, + Float16 as Float16, + Float32 as Float32, + Float64 as Float64, + Inexact as Inexact, + Int as Int, + Int4 as Int4, + Int8 as Int8, + Int16 as Int16, + Int32 as Int32, + Int64 as Int64, + Integer as Integer, + Key as Key, + Num as Num, + PRNGKeyArray as PRNGKeyArray, + Real as Real, + Scalar as Scalar, + ScalarLike as ScalarLike, + Shaped as Shaped, + UInt as UInt, + UInt4 as UInt4, + UInt8 as UInt8, + UInt16 as UInt16, + UInt32 as UInt32, + UInt64 as UInt64, + ) + + # Set up to deliberately confuse a static type checker. + PyTree: typing_extensions.TypeAlias = getattr(typing, "foo" + "bar") + # What's going on with this madness? + # + # At static-type-checking-time, we want `PyTree` to be a type for which both + # `PyTree` and `PyTree[Foo]` are equivalent to `Any`. + # (The intention is that `PyTree` be a runtime-only type; there's no real way to + # do more with static type checkers.) + # + # Unfortunately, this isn't possible: `Any` isn't subscriptable. And there's no + # equivalent way we can fake this using typing annotations. (In some sense the + # closest thing would be a `Protocol[T]` with no methods, but that's actually the + # opposite of what we want: that ends up allowing nothing at all.) + # + # The good news for us is that static type checkers have an internal escape hatch. + # If they can't figure out what a type is, then they just give up and allow + # anything. (I believe this is sometimes called `Unknown`.) Thus, this odd-looking + # annotation, which static type checkers aren't smart enough to resolve. +else: + from ._array_types import ( BFloat16 as BFloat16, Bool as Bool, Complex as Complex, @@ -96,35 +125,6 @@ if typing.TYPE_CHECKING: Real as Real, Shaped as Shaped, UInt as UInt, - Uint4 as Uint4, - UInt8 as UInt8, - UInt16 as UInt16, - UInt32 as UInt32, - UInt64 as UInt64, - ) -else: - from ._array_types import ( - BFloat16 as BFloat16, - Bool as Bool, - Complex as Complex, - Complex64 as Complex64, - Complex128 as Complex128, - Float as Float, - Float16 as Float16, - Float32 as Float32, - Float64 as Float64, - Inexact as Inexact, - Int as Int, - Int4 as Int4, - Int8 as Int8, - Int16 as Int16, - Int32 as Int32, - Int64 as Int64, - Integer as Integer, - Num as Num, - Real as Real, - Shaped as Shaped, - UInt as UInt, UInt4 as UInt4, UInt8 as UInt8, UInt16 as UInt16, @@ -132,90 +132,98 @@ else: UInt64 as UInt64, ) - if has_jax: - import jax.typing + # But crucially, does not actually import jax at all. We do that dynamically in + # __getattr__ if required. See #178. + if importlib.util.find_spec("jax") is not None: - from ._array_types import Key as Key + @ft.cache + def __getattr__(item): + if item == "Array": + if getattr(typing, "GENERATING_DOCUMENTATION", False): - if hasattr(jax.typing, "DTypeLike"): - from jax.typing import DTypeLike as DTypeLike + class Array: + pass + Array.__module__ = "builtins" + return Array + else: + import jax -# Now import PyTreeDef and PyTree -if typing.TYPE_CHECKING: - import typing_extensions + return jax.Array + elif item == "ArrayLike": + if getattr(typing, "GENERATING_DOCUMENTATION", False): - from jax.tree_util import PyTreeDef as PyTreeDef + class ArrayLike: + pass - # Set up to deliberately confuse a static type checker. - PyTree: typing_extensions.TypeAlias = getattr(typing, "foo" + "bar") - # What's going on with this madness? - # - # At static-type-checking-time, we want `PyTree` to be a type for which both - # `PyTree` and `PyTree[Foo]` are equivalent to `Any`. - # (The intention is that `PyTree` be a runtime-only type; there's no real way to - # do more with static type checkers.) - # - # Unfortunately, this isn't possible: `Any` isn't subscriptable. And there's no - # equivalent way we can fake this using typing annotations. (In some sense the - # closest thing would be a `Protocol[T]` with no methods, but that's actually the - # opposite of what we want: that ends up allowing nothing at all.) - # - # The good news for us is that static type checkers have an internal escape hatch. - # If they can't figure out what a type is, then they just give up and allow - # anything. (I believe this is sometimes called `Unknown`.) Thus, this odd-looking - # annotation, which static type checkers aren't smart enough to resolve. -elif has_jax: - if hasattr(typing, "GENERATING_DOCUMENTATION"): - # Most parts of the Equinox ecosystem have - # `typing.GENERATING_DOCUMENTATION = True` when generating documentation, to - # add whatever shims are necessary to get pretty docs. E.g. to have type - # annotations appear as just `PyTree`, not `jaxtyping.PyTree`. - # - # As jaxtyping actually wants things to appear as e.g. `jaxtyping.PyTree`, - # rather than just `PyTree`, then it sets - # `typing.GENERATING_DOCUMENTATION = False`, to disable these shims. - # - # Here we do only a `hasattr` check, as we want to get this version of - # `PyTreeDef` in both the jaxtyping and the Equinox(/etc.) docs. + ArrayLike.__module__ = "builtins" + return ArrayLike + else: + import jax.typing - class PyTreeDef: - """Alias for `jax.tree_util.PyTreeDef`, which is the type of the return - from `jax.tree_util.tree_structure(...)`. - """ + return jax.typing.ArrayLike + elif item == "PRNGKeyArray": + if getattr(typing, "GENERATING_DOCUMENTATION", False): - if typing.GENERATING_DOCUMENTATION: - # Equinox etc. docs get just `PyTreeDef`. - # jaxtyping docs get `jaxtyping.PyTreeDef`. - PyTreeDef.__module__ = "builtins" + class PRNGKeyArray: + pass - else: - from jax.tree_util import PyTreeDef as PyTreeDef + PRNGKeyArray.__module__ = "builtins" + return PRNGKeyArray + else: + # New-style `jax.random.key` have scalar shape and dtype `key`. + # Old-style `jax.random.PRNGKey` have shape `(2,)` and dtype + # `uint32`. + import jax - from ._pytree_type import PyTree as PyTree # noqa: F401 + return Union[Key[jax.Array, ""], UInt32[jax.Array, "2"]] + elif item == "DTypeLike": + import jax.typing + return jax.typing.DTypeLike + elif item == "Scalar": + import jax -# Conveniences -if typing.TYPE_CHECKING: - from ._indirection import ( - PRNGKeyArray as PRNGKeyArray, - Scalar as Scalar, - ScalarLike as ScalarLike, - ) -elif has_jax: - from ._array_types import Scalar, ScalarLike # noqa: F401 + return Shaped[jax.Array, ""] + elif item == "ScalarLike": + import jax.typing - if getattr(typing, "GENERATING_DOCUMENTATION", False): - # That is, we're generating some downstream documentation, not the jaxtyping - # documentation itself. - class PRNGKeyArray: - pass + return Shaped[jax.typing.ArrayLike, ""] + elif item == "PyTree": + from ._pytree_type import PyTree - PRNGKeyArray.__module__ = "builtins" - else: - from ._array_types import PRNGKeyArray + return PyTree + elif item == "PyTreeDef": + if hasattr(typing, "GENERATING_DOCUMENTATION"): + # Most parts of the Equinox ecosystem have + # `typing.GENERATING_DOCUMENTATION = True` when generating + # documentation, to add whatever shims are necessary to get pretty + # docs. E.g. to have type annotations appear as just `PyTree`, not + # `jaxtyping.PyTree`. + # + # As jaxtyping actually wants things to appear as e.g. + # `jaxtyping.PyTree`, rather than just `PyTree`, then it sets + # `typing.GENERATING_DOCUMENTATION = False`, to disable these shims. + # + # Here we do only a `hasattr` check, as we want to get this version + # of `PyTreeDef` in both the jaxtyping and the Equinox(/etc.) docs. -del has_jax + class PyTreeDef: + """Alias for `jax.tree_util.PyTreeDef`, which is the type of the + return from `jax.tree_util.tree_structure(...)`. + """ + + if typing.GENERATING_DOCUMENTATION: + # Equinox etc. docs get just `PyTreeDef`. + # jaxtyping docs get `jaxtyping.PyTreeDef`. + PyTreeDef.__module__ = "builtins" + return PyTreeDef + else: + import jax.tree_util + + return jax.tree_util.PyTreeDef + else: + raise AttributeError(f"module jaxtyping has no attribute {item!r}") check_equinox_version = True # easy-to-replace line with copybara diff --git a/jaxtyping/_array_types.py b/jaxtyping/_array_types.py index d1049af..e5eaaff 100644 --- a/jaxtyping/_array_types.py +++ b/jaxtyping/_array_types.py @@ -36,18 +36,6 @@ from ._storage import ( ) -try: - import jax -except (ImportError, RuntimeError, AttributeError): - # We catch `RuntimeError` as JAX will throw this if it's present, but unable to run - # on the current machine. This fails with this error. - # We catch `AttributeError` as the above then leaves the module in a partially - # initialised state, which causes subsequent imports to fail with this error. - has_jax = False -else: - has_jax = True - - _array_name_format = "dtype_and_shape" @@ -721,10 +709,4 @@ Num = _make_dtype(uints + ints + floats + complexes, "Num") Shaped = _make_dtype(_any_dtype, "Shaped") -if has_jax: - Key = _make_dtype(_prng_key, "Key") - # New-style `jax.random.key` have scalar shape and dtype `key`. - # Old-style `jax.random.PRNGKey` have shape `(2,)` and dtype `uint32`. - PRNGKeyArray = Union[Key[jax.Array, ""], UInt32[jax.Array, "2"]] - Scalar = Shaped[jax.Array, ""] - ScalarLike = Shaped[jax.typing.ArrayLike, ""] +Key = _make_dtype(_prng_key, "Key") diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py index 9a09861..d23ad16 100644 --- a/jaxtyping/_decorator.py +++ b/jaxtyping/_decorator.py @@ -19,21 +19,13 @@ import dataclasses import functools as ft +import importlib.util import inspect import itertools as it import sys import warnings from typing import Any, get_args, get_origin, get_type_hints, overload - -try: - import jax._src.traceback_util as traceback_util -except ImportError: - pass -else: - traceback_util.register_exclusion(__file__) - - from ._config import config from ._errors import AnnotationError, TypeCheckError from ._storage import pop_shape_memo, push_shape_memo, shape_str @@ -45,6 +37,7 @@ class _Sentinel: _sentinel = _Sentinel() +_tb_flag = True @overload @@ -193,6 +186,13 @@ def jaxtyped(fn=_sentinel, *, typechecker=_sentinel): useful when working at the global scope. """ + global _tb_flag + if _tb_flag and importlib.util.find_spec("jax._src.traceback_util") is not None: + import jax._src.traceback_util as traceback_util + + traceback_util.register_exclusion(__file__) + _tb_flag = False + # First handle the `jaxtyped("context")` usage, which is a special case. if fn == "context": if typechecker is not _sentinel: diff --git a/jaxtyping/_indirection.py b/jaxtyping/_indirection.py index ab70b83..5660b68 100644 --- a/jaxtyping/_indirection.py +++ b/jaxtyping/_indirection.py @@ -48,8 +48,13 @@ from typing import ( Annotated as UInt16, # noqa: F401 Annotated as UInt32, # noqa: F401 Annotated as UInt64, # noqa: F401 + TYPE_CHECKING, ) + +if not TYPE_CHECKING: + assert False + from jax import ( Array as PRNGKeyArray, # noqa: F401 Array as Scalar, # noqa: F401 diff --git a/pyproject.toml b/pyproject.toml index 6caabdd..b75ad2f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Mathematics", ] urls = {repository = "https://github.com/google/jaxtyping" } -dependencies = ["numpy>=1.20.0", "typeguard>=2.13.3,<3", "typing_extensions>=3.7.4.1"] +dependencies = ["numpy>=1.20.0", "typeguard==2.13.3"] entry-points = {pytest11 = {jaxtyping = "jaxtyping._pytest_plugin"}} [build-system] diff --git a/test/test_all_importable.py b/test/test_all_importable.py new file mode 100644 index 0000000..2fead03 --- /dev/null +++ b/test/test_all_importable.py @@ -0,0 +1,47 @@ +# We have some pretty complicated semantics in `__init__.py`. +# Here we check that we didn't miss one of them on our runtime branch. +def test_all_importable(): + # Ordered according to their appearance in the documentation. + from jaxtyping import ( # noqa: I001 + Shaped, # noqa: F401 + Bool, # noqa: F401 + Key, # noqa: F401 + Num, # noqa: F401 + Inexact, # noqa: F401 + Float, # noqa: F401 + BFloat16, # noqa: F401 + Float16, # noqa: F401 + Float32, # noqa: F401 + Float64, # noqa: F401 + Complex, # noqa: F401 + Complex64, # noqa: F401 + Complex128, # noqa: F401 + Integer, # noqa: F401 + UInt, # noqa: F401 + UInt4, # noqa: F401 + UInt8, # noqa: F401 + UInt16, # noqa: F401 + UInt32, # noqa: F401 + UInt64, # noqa: F401 + Int, # noqa: F401 + Int4, # noqa: F401 + Int8, # noqa: F401 + Int16, # noqa: F401 + Int32, # noqa: F401 + Int64, # noqa: F401 + Real, # noqa: F401 + Array, # noqa: F401 + ArrayLike, # noqa: F401 + Scalar, # noqa: F401 + ScalarLike, # noqa: F401 + PRNGKeyArray, # noqa: F401 + PyTreeDef, # noqa: F401 + PyTree, # noqa: F401 + jaxtyped, # noqa: F401 + install_import_hook, # noqa: F401 + AbstractArray, # noqa: F401 + AbstractDtype, # noqa: F401 + print_bindings, # noqa: F401 + get_array_name_format, # noqa: F401 + set_array_name_format, # noqa: F401 + ) diff --git a/test/test_no_jax_dependency.py b/test/test_no_jax_dependency.py new file mode 100644 index 0000000..7b226cd --- /dev/null +++ b/test/test_no_jax_dependency.py @@ -0,0 +1,20 @@ +import subprocess + + +def test_no_jax_dependency(): + result = subprocess.run( + "python -c 'import jaxtyping; import sys; sys.exit(\"jax\" in sys.modules)'", + shell=True, + ) + assert result.returncode == 0 + + +# Meta-test: test that the above test will work. (i.e. that I haven't messed up using +# subprocess.) +def test_meta(): + result = subprocess.run( + "python -c 'import jaxtyping; import jax; import sys; " + 'sys.exit("jax" in sys.modules)\'', + shell=True, + ) + assert result.returncode == 1