Skip to main content

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

Project description

jaxtyping

Type annotations and runtime type-checking for:

  1. shape and dtype of JAX arrays; (Now also supports PyTorch, NumPy, and TensorFlow!)
  2. PyTrees.

For example:

from jaxtyping import Array, Float, PyTree

# Accepts floating-point 2D arrays with matching axes
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 Python 3.9+.

JAX is an optional dependency, required for a few JAX-specific types. If JAX is not installed then these will not be available, but you may still use jaxtyping to provide shape/dtype annotations for PyTorch/NumPy/TensorFlow/etc.

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).

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.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

jaxtyping-0.2.35.tar.gz (45.1 kB view details)

Uploaded Source

Built Distribution

jaxtyping-0.2.35-py3-none-any.whl (55.8 kB view details)

Uploaded Python 3

File details

Details for the file jaxtyping-0.2.35.tar.gz.

File metadata

  • Download URL: jaxtyping-0.2.35.tar.gz
  • Upload date:
  • Size: 45.1 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.0 CPython/3.11.9

File hashes

Hashes for jaxtyping-0.2.35.tar.gz
Algorithm Hash digest
SHA256 86095171e33cb76290f704706ffc70297934ec1759422a7bc8e5918ad9072fb6
MD5 8cdbe398eafc6cd3c224dd079d4299de
BLAKE2b-256 244161a359b5a05b8016429011652095b56b87a1ce03fc907575c509ded220bb

See more details on using hashes here.

File details

Details for the file jaxtyping-0.2.35-py3-none-any.whl.

File metadata

  • Download URL: jaxtyping-0.2.35-py3-none-any.whl
  • Upload date:
  • Size: 55.8 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.0 CPython/3.11.9

File hashes

Hashes for jaxtyping-0.2.35-py3-none-any.whl
Algorithm Hash digest
SHA256 70c13a622e0e42b9b60c62c21a6cf0ca3480a5ab77071ba39865fd778840029e
MD5 6e5a42e8ac56d898bf367a3c3814daa8
BLAKE2b-256 2019bb407b6bfaa46c312033dad1f8b44cdac4c848a8852d125ea514f72860db

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page