Skip to content

Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees.

License

Notifications You must be signed in to change notification settings

murphyk/jaxtyping

 
 

Repository files navigation

jaxtyping

Type annotations and runtime checking for:

  1. shape and dtype of JAX arrays;
  2. PyTrees.

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

Full API reference

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.

About

Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees.

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Python 100.0%