wassname/jaxtyping
1
0
Fork 0
mirror of https://github.com/wassname/jaxtyping.git synced 2026-09-11 12:21:38 +08:00
Code Issues Packages Projects Releases Wiki Activity

No longer imports JAX at all! This is done dynamically when required. See #178

Browse Source
This commit is contained in:
Patrick Kidger
2024-02-25 12:07:01 +00:00
parent 9beb5f2d29
commit 17ea4b13eb
7 changed files with 219 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files
jaxtyping/__init__.py
+136 -128
View File
@@ -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<foo>`.
# 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
jaxtyping/_array_types.py
+1 -19
View File
@@ -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<foo>`.
# 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")
jaxtyping/_decorator.py
+9 -9
View File
@@ -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:
jaxtyping/_indirection.py
+5
View File
@@ -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
pyproject.toml
+1 -1
View File
@@ -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]
test/test_all_importable.py
+47
View File
@@ -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
)
test/test_no_jax_dependency.py
+20
View File
@@ -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