---
id: jaxtyping
version: "0.3.11"
license: MIT License Copyright (c) 2022 Google LLC Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in… (full text in the JSON record)
license_treatment: permissive
maintenance: active
---

# jaxtyping — Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays.

License: permissive · Maintenance: active · Downloads: 7.2M/mo

## 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 above — 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

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

- License: MIT License Copyright (c) 2022 Google LLC Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in… (full text in the JSON record) (permissive)
- Python support: supports_current
- Install friction: low
- Maintenance: active
- Downloads: 7.2M/month (top 5,000 on PyPI)
- Known vulnerabilities: none known

## Tags

array shape type annotations, tensor dtype checking, runtime array validation, numpy pytorch jax typing, shape and dtype annotations, tensor type hints, array dimension validation, type-checking, array-validation, deep-learning

[View on SkillFed](https://skillfed.io/packages/jaxtyping) · [View on PyPI](https://pypi.org/project/jaxtyping/)