Roman Knyazhitskiy
172b83b4fc
* Add a test for generators * Remove output annotations from decorators Also guarded torch imports for better compatibility with requirements.txt * Add flag to the main meta class to skip the typecheck * Return to the old solution * Make async tests work * Minor adjustments/fixing typos * Correct Python path for new tests * Remove some jax-dependent code * Implement equality for MetaArrays * Make all Dim variations frozen dataclasses * Shorten AbstractArray methods * Final touches * Removing get_origin use * Update tests with @jaxtyp
jaxtyping
Type annotations and runtime type-checking for:
For example:
from jaxtyping import Array, Float, PyTree
# Accepts floating-point 2D arrays with matching axes
def matrix_multiply(x: Float[Array, "dim1 dim2"],
y: Float[Array, "dim2 dim3"]
) -> Float[Array, "dim1 dim3"]:
...
def accepts_pytree_of_ints(x: PyTree[int]):
...
def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
...
Installation
pip install jaxtyping
Requires Python 3.9+.
JAX is an optional dependency, required for a few JAX-specific types. If JAX is not installed then these will not be available, but you may still use jaxtyping to provide shape/dtype annotations for PyTorch/NumPy/TensorFlow/etc.
The annotations provided by jaxtyping are compatible with runtime type-checking packages, so it is common to also install one of these. The two most popular are typeguard (which exhaustively checks every argument) and beartype (which checks random pieces of arguments).
Documentation
Available at https://docs.kidger.site/jaxtyping.
Finally
See also: other libraries in the JAX ecosystem
Equinox: neural networks.
Optax: first-order gradient (SGD, Adam, ...) optimisers.
Diffrax: numerical differential equation solvers.
Optimistix: root finding, minimisation, fixed points, and least squares.
Lineax: linear solvers.
BlackJAX: probabilistic+Bayesian sampling.
Orbax: checkpointing (async/multi-host/multi-device).
sympy2jax: SymPy<->JAX conversion; train symbolic expressions via gradient descent.
Eqxvision: computer vision models.
Levanter: scalable+reliable training of foundation models (e.g. LLMs).
PySR: symbolic regression. (Non-JAX honourable mention!)
Disclaimer
This is not an official Google product.