Patrick Kidger
12d540794f
Phew, this ended up being a pretty complicated change! The basic summary is that we now support the syntax ``` @jaxtyped(typechecker=typechecker) def f(...): ... ``` and when using this, we now get pretty error messages about what went wrong. ( The old syntax, i.e. ``` @jaxtyped @typechecker def f(...): ... ``` is still supported, but doesn't give much information. ) The internals of this do quite a lot of magic! In particular we dynamically create quite a lot of functions and test the provided arguments against their signatures. The overhead should still be minimal under `jax.jit`, though. (TODO: what's the overhead like in non-jit situations, e.g. PyTorch? I've tried to minimise the overhead throughout just to be sure, but perhaps PyTorch users should stick to the old syntax?)
2.3 KiB
Runtime type checking
(See the FAQ for details on static type checking.)
Runtime type checking synergises beautifully with jax.jit! All shape checks will be performed only whilst tracing, and will not impact runtime performance.
There are two approaches: either use [jaxtyping.jaxtyped][] to typecheck a single function, or [jaxtyping.install_import_hook][] to typecheck a whole codebase.
In either case, the actual business of checking types is performed with the help of a runtime type-checking library. The two most popular are beartype and typeguard. (If using typeguard, then specifically the version 2.* series should be used. Later versions -- 3 and 4 -- have some known issues.)
::: jaxtyping.jaxtyped
::: jaxtyping.install_import_hook
Pytest hook
The import hook can be installed at test-time only, as a pytest hook. From the command line the syntax is:
pytest --jaxtyping-packages=foo,bar.baz,beartype.beartype
or in pyproject.toml:
[tool.pytest.ini_options]
addopts = "--jaxtyping-packages=foo,bar.baz,beartype.beartype"
or in pytest.ini:
[pytest]
addopts = --jaxtyping-packages=foo,bar.baz,beartype.beartype
This example will apply the import hook to all modules whose names start with either foo or bar.baz. The typechecker used in this example is beartype.beartype.
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:
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.
Other runtime type-checking libraries
Beartype and typeguard happen to be the two most popular runtime type-checking libraries (at least at time of writing), but jaxtyping should be compatible with all runtime type checkers out-of-the-box. The runtime type-checking library just needs to provide a type-checking decorator (analgous to beartype.beartype or typeguard.typechecked), and perform isinstance checks against jaxtyping's types.