Type annotations and runtime checking for:
For example:
from jaxtyping import f32, PyTree
def matrix_multiply(x: f32["dim1 dim2"], y: f32["dim2 dim3"]) -> f32["dim1 dim3"]:
...
def accepts_pytree_of_ints(x: PyTree[int]):
...
def accepts_pytree_of_arrays(x: PyTree[f32["batch c1 c2"]]):
...
pip install jaxtyping
Requires JAX 0.3.4+.
Also install your favourite runtime type-checking package. The two most popular are typeguard (which exhaustively checks every argument) and beartype (which checks random pieces of arguments).
FAQ (static type checking, flake8, etc.)
Neural networks: Equinox.
Numerical differential equation solvers: Diffrax.
SymPy<->JAX conversion; train symbolic expressions via gradient descent: sympy2jax.
Shape annotations + runtime type checking is inspired by TorchTyping.
The concise syntax is inspired by etils.array_types.
This is not an official Google product.