jaxtyping (pypi): Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees.
Keywords: deep-learning|equinox|jax|neural-networks|typing