Document IPython extension and version nump

README.md

@@ -41,15 +41,21 @@ Available at [https://docs.kidger.site/jaxtyping](https://docs.kidger.site/jaxty
 
 ## Finally
 
-### See also: other tools in the JAX ecosystem
+### See also: other libraries in the JAX ecosystem
 
-Neural networks: [Equinox](https://github.com/patrick-kidger/equinox).
+[Equinox](https://github.com/patrick-kidger/equinox): neural networks.
 
-Numerical differential equation solvers: [Diffrax](https://github.com/patrick-kidger/diffrax).
+[Optax](https://github.com/deepmind/optax): first-order gradient (SGD, Adam, ...) optimisers.
 
-Computer vision models: [Eqxvision](https://github.com/paganpasta/eqxvision).
+[Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.
 
-SymPy<->JAX conversion; train symbolic expressions via gradient descent: [sympy2jax](https://github.com/google/sympy2jax).
+[Lineax](https://github.com/google/lineax): linear solvers and linear least squares.
+
+[Eqxvision](https://github.com/paganpasta/eqxvision): computer vision models.
+
+[sympy2jax](https://github.com/google/sympy2jax): SymPy<->JAX conversion; train symbolic expressions via gradient descent.
+
+[Levanter](https://github.com/stanford-crfm/levanter): scalable+reliable training of foundation models (e.g. LLMs).
 
 ### Disclaimer
 

docs/api/array.md

@@ -95,16 +95,15 @@ BatchImage = Float[Array, "batch channels height width"]
 
 Note that `jaxtyping.{Array, ArrayLike}` are only available if JAX has been installed.
 
-## Scalars, PRNGKeys
+## Scalars, PRNG keys
 
 For convenience, jaxtyping also includes `jaxtyping.Scalar`, `jaxtyping.ScalarLike`, and `jaxtyping.PRNGKeyArray`, defined as:
 ```python
 Scalar = Shaped[Array, ""]
 ScalarLike = Shaped[ArrayLike, ""]
 
-# Depending on the value of `JAX_ENABLE_CUSTOM_PRNG`:
+# Left: new-style typed keys; right: old-style keys. See JEP 9263.
-PRNGKeyArray = Key[Array, ""]
+PRNGKeyArray = Union[Key[Array, ""], UInt32[Array, "2"]]
-PRNGKeyArray = UInt32[Array, "2"]
 ```
 
 Recalling that shape-and-dtype specified jaxtyping arrays can be nested, this means that e.g. you can annotate the output of `jax.random.split` with `Shaped[PRNGKeyArray, "2"]`, or e.g. an integer scalar with `Int[Scalar, ""]`.

docs/api/runtime-type-checking.md

@@ -20,3 +20,15 @@ It can be a lot of effort to add `@jaxtyped` decorators all over your codebase.
 The easier option is usually to use the import hook.
 
 ::: jaxtyping.install_import_hook
+
+---
+
+#### IPython extension
+
+If you are running in an IPython environment (for example a Jupyter or Colab notebook), then the jaxtyping hook can be automatically ran via a custom magic:
+```python
+import jaxtyping
+%load_ext jaxtyping
+%jaxtyping.typechecker beartype.beartype  # or any other runtime type checker
+```
+Place this at the start of your notebook -- everything that is directly defined in the notebook, after this magic is run, will be hook'd.

docs/index.md

@@ -40,3 +40,19 @@ def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
 ## Next steps
 
 Have a read of the [Array annotations](./api/array.md) documentation on the left-hand bar!
+
+## See also: other libraries in the JAX ecosystem
+
+[Equinox](https://github.com/patrick-kidger/equinox): neural networks.
+
+[Optax](https://github.com/deepmind/optax): first-order gradient (SGD, Adam, ...) optimisers.
+
+[Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.
+
+[Lineax](https://github.com/google/lineax): linear solvers and linear least squares.
+
+[Eqxvision](https://github.com/paganpasta/eqxvision): computer vision models.
+
+[sympy2jax](https://github.com/google/sympy2jax): SymPy<->JAX conversion; train symbolic expressions via gradient descent.
+
+[Levanter](https://github.com/stanford-crfm/levanter): scalable+reliable training of foundation models (e.g. LLMs).

jaxtyping/__init__.py

@@ -20,24 +20,17 @@
 import importlib.metadata
 import typing
 
-
-try:
-    import jax
-except ImportError:
-    has_jax = False
-else:
-    has_jax = True
-    del jax
-
 # 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 ._decorator import jaxtyped as jaxtyped
 from ._import_hook import install_import_hook as install_import_hook
+from ._ipython_extension import load_ipython_extension as load_ipython_extension
 
 
 # Now import Array and ArrayLike
@@ -154,12 +147,28 @@ if typing.TYPE_CHECKING:
     # 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.
 
         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"
+
     else:
         from jax.tree_util import PyTreeDef as PyTreeDef
 
@@ -172,7 +181,17 @@ if typing.TYPE_CHECKING:
 
     from ._indirection import Scalar as Scalar, ScalarLike as ScalarLike
 elif has_jax:
-    from ._array_types import PRNGKeyArray, Scalar, ScalarLike  # noqa: F401
+    from ._array_types import Scalar, ScalarLike  # noqa: F401
+
+    if getattr(typing, "GENERATING_DOCUMENTATION", False):
+        # That is, we're generating some downstream documentation, not the jaxtyping
+        # documentation itself.
+        class PRNGKeyArray:
+            pass
+
+        PRNGKeyArray.__module__ = "builtins"
+    else:
+        from ._array_types import PRNGKeyArray
 
 del has_jax
 

jaxtyping/_array_types.py

@@ -32,7 +32,11 @@ from ._decorator import storage
 
 try:
     import jax
-except ImportError:
+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
@@ -137,12 +141,21 @@ def _check_dims(
     return True
 
 
+def _is_jax_extended_dtype(dtype: Any) -> bool:
+    if not has_jax:
+        return False
+    if hasattr(jax.dtypes, "extended"):  # jax>=0.4.14
+        return jax.numpy.issubdtype(dtype, jax.dtypes.extended)
+    else:  # jax<=0.4.13
+        return jax.core.is_opaque_dtype(dtype)
+
+
 class _MetaAbstractArray(type):
     def __instancecheck__(cls, obj):
         if not isinstance(obj, cls.array_type):
             return False
 
-        if has_jax and jax.core.is_opaque_dtype(obj.dtype):
+        if _is_jax_extended_dtype(obj.dtype):
             dtype = str(obj.dtype)
         elif hasattr(obj.dtype, "type") and hasattr(obj.dtype.type, "__name__"):
             # JAX, numpy
@@ -512,8 +525,9 @@ class _MetaAbstractDtype(type):
     def __getitem__(cls, item: tuple[Any, str]):
         if not isinstance(item, tuple) or len(item) != 2:
             raise ValueError(
-                "As of jaxtyping v0.2.0, type annotations must now include an explicit "
+                "As of jaxtyping v0.2.0, type annotations must now include both an "
-                "array type. For example `jaxtyping.Float32[jax.Array, 'foo bar']`."
+                "array type and a shape. For example `Float[Array, 'foo bar']`.\n"
+                "Ellipsis can be used to accept any shape: `Float[Array, '...']`."
             )
         array_type, dim_str = item
         del item
@@ -648,12 +662,10 @@ Num = _make_dtype(uints + ints + floats + complexes, "Num")
 Shaped = _make_dtype(_any_dtype, "Shaped")
 
 if has_jax:
-    if jax.config.jax_enable_custom_prng:
+    _key_regex = re.compile(r"^key<\w+>$")
-        _key_regex = re.compile(r"^key<\w+>$")
+    Key = _make_dtype(_key_regex, "Key")
-        Key = _make_dtype(_key_regex, "Key")
+    # New-style `jax.random.key` have scalar shape and dtype `key<foo>`.
-        PRNGKeyArray = Key[jax.Array, ""]
+    # Old-style `jax.random.PRNGKey` have shape `(2,)` and dtype `uint32`.
-    else:
+    PRNGKeyArray = Union[Key[jax.Array, ""], UInt32[jax.Array, "2"]]
-        Key = UInt32
-        PRNGKeyArray = Key[jax.Array, "2"]
     Scalar = Shaped[jax.Array, ""]
     ScalarLike = Shaped[jax.typing.ArrayLike, ""]

jaxtyping/_decorator.py

@@ -25,6 +25,14 @@ import types
 import weakref
 
 
+try:
+    import jax._src.traceback_util as traceback_util
+except ImportError:
+    pass
+else:
+    traceback_util.register_exclusion(__file__)
+
+
 storage = threading.local()
 
 

jaxtyping/_import_hook.py

@@ -51,6 +51,7 @@
 
 import ast
 import functools as ft
+import hashlib
 import sys
 from collections.abc import Sequence
 from importlib.abc import MetaPathFinder
@@ -165,7 +166,9 @@ class _JaxtypingLoader(SourceFileLoader):
     def __init__(self, *args, typechecker, **kwargs):
         super().__init__(*args, **kwargs)
         self._typechecker = typechecker
-        self._typechecker_hash = str(abs(hash(self._typechecker)))
+        self._typechecker_hash = hashlib.md5(
+            self._typechecker.encode("utf-8")
+        ).hexdigest()
 
     def source_to_code(self, data, path, *, _optimize=-1):
         source = decode_source(data)

jaxtyping/_ipython_extension.py

@@ -0,0 +1,35 @@
+from ._import_hook import _JaxtypingTransformer
+
+
+try:
+    from IPython.core.magic import line_magic, Magics, magics_class
+
+    @magics_class
+    class ChooseTypecheckerMagics(Magics):
+        @line_magic("jaxtyping.typechecker")
+        def typechecker(self, typechecker):
+            # remove old _JaxtypingTransformer, if present
+            self.shell.ast_transformers = list(
+                filter(
+                    lambda x: not isinstance(x, _JaxtypingTransformer),
+                    self.shell.ast_transformers,
+                )
+            )
+
+            # add new one
+            self.shell.ast_transformers.append(
+                _JaxtypingTransformer(typechecker=typechecker)
+            )
+
+except ImportError:
+    pass
+
+
+def load_ipython_extension(ipython):
+    try:
+        ipython.register_magics(ChooseTypecheckerMagics)
+    except NameError:
+        raise NameError(
+            "ChooseTypecheckerMagics is not defined.\n\n"
+            + "You may be trying to use IPython extension without IPython installed."
+        )

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jaxtyping"
-version = "0.2.20"
+version = "0.2.22"
 description = "Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees."
 readme = "README.md"
 requires-python ="~=3.9"

test/requirements.txt

@@ -4,3 +4,4 @@ equinox
 jaxlib
 pytest
 typeguard<3
+IPython

test/test_ipython_extension.py

@@ -0,0 +1,151 @@
+import pytest
+from IPython.testing.globalipapp import start_ipython
+
+from .helpers import ParamError
+
+
+@pytest.fixture(scope="session")
+def session_ip():
+    yield start_ipython()
+
+
+@pytest.fixture(scope="function")
+def ip(session_ip):
+    session_ip.run_cell(raw_cell="import jaxtyping")
+    session_ip.run_line_magic(magic_name="load_ext", line="jaxtyping")
+    session_ip.run_line_magic(
+        magic_name="jaxtyping.typechecker", line="beartype.beartype"
+    )
+    yield session_ip
+
+
+def test_that_ipython_works(ip):
+    ip.run_cell(raw_cell="x = 1").raise_error()
+    assert ip.user_global_ns["x"] == 1
+
+
+def test_function_beartype(ip):
+    ip.run_cell(
+        raw_cell="""
+    def f(x: int):
+        pass
+                """
+    ).raise_error()
+    ip.run_cell(raw_cell="f(1)").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell='f("x")').raise_error()
+
+
+def test_function_none(ip):
+    ip.run_cell(
+        raw_cell="""
+    def f(a,b,c):
+        pass
+                """
+    ).raise_error()
+    ip.run_cell(raw_cell='f(1,2,"k")').raise_error()
+
+
+def test_function_jaxtyped(ip):
+    ip.run_cell(
+        raw_cell="""
+    from jaxtyping import Float, Array, Int
+    import jax
+
+    def g(x: Float[Array, "1"]):
+        return x + 1
+
+                """
+    ).raise_error()
+
+    ip.run_cell(raw_cell="g(jax.numpy.array([1.0]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array(1.0))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array([1]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array([2, 3]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell='g("string")').raise_error()
+
+
+def test_function_jaxtyped_and_jitted(ip):
+    ip.run_cell(
+        raw_cell="""
+    from jaxtyping import Float, Array, Int
+    import jax
+
+    @jax.jit
+    def g(x: Float[Array, "1"]):
+        return x + 1
+
+                """
+    ).raise_error()
+
+    ip.run_cell(raw_cell="g(jax.numpy.array([1.0]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array(1.0))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array([1]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array([2, 3]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell='g("string")').raise_error()
+
+
+def test_class_jaxtyped(ip):
+    ip.run_cell(
+        raw_cell="""
+    from jaxtyping import Float, Array, Int
+    import equinox as eqx
+    import jax
+
+    class A(eqx.Module):
+        x: Float[Array, "2"]
+
+        def do_something(self, y: Int[Array, ""]):
+            return self.x + y
+                """
+    ).raise_error()
+
+    ip.run_cell(raw_cell="a = A(jax.numpy.array([1.0, 2.0]))").raise_error()
+    ip.run_cell(raw_cell="a.do_something(jax.numpy.array(2))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="A(jax.numpy.array([1.0]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(
+            raw_cell="a.do_something(jax.numpy.array([2.0, 3.0]))"
+        ).raise_error()
+
+
+def test_class_not_dataclass(ip):
+    ip.run_cell(
+        raw_cell="""
+    from jaxtyping import Float, Array, Int
+    import equinox as eqx
+    import jax
+
+    class A:
+        def __init__(self, x):
+            self.x = x
+
+        def do_something(self, y):
+            return x + y
+                """
+    ).raise_error()
+
+    ip.run_cell(raw_cell="a = A(jax.numpy.array([1.0, 2.0]))").raise_error()
+    ip.run_cell(raw_cell="a.do_something(jax.numpy.array(2))").raise_error()
+    ip.run_cell(raw_cell="A(jax.numpy.array([1.0]))").raise_error()
+    ip.run_cell(raw_cell="a.do_something(jax.numpy.array([2.0, 3.0]))").raise_error()