jaxtyping

What it is and what it does

jaxtyping is a type-annotation library that lets you specify the shape and data type of arrays and tensors directly in function signatures. Instead of writing generic array or tensor type hints, you can declare that a function expects a floating-point array with specific axis names—and optionally enforce those constraints at runtime. The library works with JAX, PyTorch, NumPy, MLX, and TensorFlow, despite its historical name.

The annotations themselves are static (compatible with standard Python type checkers), but jaxtyping is designed to pair with runtime type-checking libraries, which can then validate that actual arguments match the declared shapes and dtypes. This catches shape mismatches and dtype errors early, which is especially valuable in numerical and deep-learning code where silent broadcasting or type coercion can hide bugs.

Use it for:

• Annotate neural network layer inputs and outputs with expected tensor shapes to catch dimension mismatches early
• Document and enforce dtype constraints in scientific computing functions
• Pair with runtime type-checking tools to add validation to model code
• Improve IDE autocomplete and static type-checker support for array-heavy codebases
• Validate matrix operation arguments for compatible dimensions

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Provides type annotations and runtime type-checking for array shape and dtype across JAX, PyTorch, NumPy, MLX, and TensorFlow, with no JAX dependency required.

Yes. Active maintenance, permissive MIT license, low install friction, and no known vulnerabilities. Valuable for any codebase mixing arrays or tensors with type hints. Install it for annotations alone; add a runtime type-checking package separately if runtime validation is needed.

Install

jaxtyping on PyPI

pip install jaxtyping

uv add jaxtyping

poetry add jaxtyping

Installing jaxtyping

Before you install

Low friction install with a single runtime dependency (wadler-lindig). Active maintenance with recent releases; last commit 2026-07-08. Requires Python 3.11+.

License in practice

MIT License (permissive). Code includes sections modified from typeguard under MIT terms. Allows commercial and private use with attribution.

Quickstart

pip install jaxtyping

from jaxtyping import Float

def matrix_multiply(x: Float["dim1 dim2"],
                    y: Float["dim2 dim3"]
                  ) -> Float["dim1 dim3"]:
    pass

Requires Python 3.11+. Runtime type-checking typically requires installing a separate type-checking package.

Verify before relying

• Performance overhead of runtime type-checking on large arrays or frequent function calls
• Compatibility with type-checking tools beyond those mentioned in documentation
• Support for custom array types or frameworks beyond those listed

Package facts

Evidence: jaxtyping-0.3.11-py3-none-any.whl

Keywords: deep-learning, equinox, jax, neural-networks, typing

Tags

More Artificial Intelligence packages