jaxtyping
A library providing type annotations and runtime type-checking for the shape and dtype of JAX/PyTorch/NumPy/MLX/TensorFlow arrays and tensors.
The name 'jax'typing is now historical, we support all of the above and have no JAX dependency!
from jaxtyping import Float
from torch import Tensor
# Accepts floating-point 2D arrays with matching axes
def matrix_multiply(x: Float[Tensor, "dim1 dim2"],
y: Float[Tensor, "dim2 dim3"]
) -> Float[Tensor, "dim1 dim3"]:
...
Installation
pip install jaxtyping
Requires Python 3.10+.
The annotations provided by jaxtyping are compatible with runtime type-checking packages, so it is common to also install one of these. The two most popular are typeguard (which exhaustively checks every argument) and beartype (which checks random pieces of arguments).
Documentation
Available at https://docs.kidger.site/jaxtyping.
See also: other libraries in the JAX ecosystem
Always useful
Equinox: neural networks and everything not already in core JAX!
Deep learning
Optax: first-order gradient (SGD, Adam, ...) optimisers.
Orbax: checkpointing (async/multi-host/multi-device).
Levanter: scalable+reliable training of foundation models (e.g. LLMs).
paramax: parameterizations and constraints for PyTrees.
Scientific computing
Diffrax: numerical differential equation solvers.
Optimistix: root finding, minimisation, fixed points, and least squares.
Lineax: linear solvers.
BlackJAX: probabilistic+Bayesian sampling.
sympy2jax: SymPy<->JAX conversion; train symbolic expressions via gradient descent.
PySR: symbolic regression. (Non-JAX honourable mention!)
Awesome JAX
Awesome JAX: a longer list of other JAX projects.
1 maintainer
Metadata
1 maintainer
Release files for jaxtyping 0.3.11
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| jaxtyping-0.3.11.tar.gz | 46.5 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| jaxtyping-0.3.11-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 103.1 kB
Release files / jaxtyping-0.3.11.tar.gz
| Download URL | jaxtyping-0.3.11.tar.gz |
|---|---|
| Size | 46.5 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
b09c14acf6686feb9e0df5b0d8c6e7c5b6f8d36bf059ee54cd522a186c2ef050
|
|
BLAKE2b-256 checksum How to use checksums |
91c1091b8852bd7cbf50bd655543c8506033cf4029300c67f8c176c1286879a9
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.2.0 CPython/3.12.13
|
Release files / jaxtyping-0.3.11-py3-none-any.whl
| Download URL | jaxtyping-0.3.11-py3-none-any.whl |
|---|---|
| Size | 56.6 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
8a4bedc4e3f963fa82df41bd13c7ebc2bad925601eb48614c65798f21329d4e3
|
|
BLAKE2b-256 checksum How to use checksums |
8b38c66bbdc5047f4776c2bd3e47e5295a350e3fa44d5b8942105e71c2a876a0
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.2.0 CPython/3.12.13
|
