Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees.
conda install conda-forge::jaxtyping