Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees.
Project description
jaxtyping
Type annotations and runtime checking for:
For example:
from jaxtyping import Array, Float, PyTree
# Accepts floating-point 2D arrays with matching dimensions
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 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).
Documentation
FAQ (static type checking, flake8, etc.)
Finally
See also: other tools in the JAX ecosystem
Neural networks: Equinox.
Numerical differential equation solvers: Diffrax.
SymPy<->JAX conversion; train symbolic expressions via gradient descent: sympy2jax.
Acknowledgements
Shape annotations + runtime type checking is inspired by TorchTyping.
The concise syntax is partially inspired by etils.array_types.
Disclaimer
This is not an official Google product.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for jaxtyping-0.2.6-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 41cc38ae665ac557196d56e2f271e9a95010fc5953cb5cdbeeed2bed5bf08d78 |
|
MD5 | a825a2daf37ee8a96218355800ab06d9 |
|
BLAKE2b-256 | 3f98578c1459c56d5656dfa8c93e7f03d37b40389eb077ffd6ff9539d5d2fa87 |