jaxtyping
Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays.
Decision gist · record as of 2026-08-14
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.AI-flagged interpretation of the facts on this page — verify before relying
Before you install
- Requires Python 3.11+.
- Runtime type-checking typically requires installing a separate type-checking package.
- Low friction install with a single runtime dependency (wadler-lindig).
License · maintenance · safety
permissive license (permissive) — MIT License (permissive). Code includes sections modified from typeguard under MIT terms. Allows commercial and private use with attribution.
last release 2026-06-13 (62 days) · last repo commit 2026-07-08 · 1,854 stars
0 known vulnerabilities (OSV.dev, 2026-08-14) · 7,153,117 downloads/mo, #1,779 on PyPI
Alternatives
Verify before relying
pip install jaxtyping
from jaxtyping import Float
def matrix_multiply(x: Float["dim1 dim2"],
y: Float["dim2 dim3"]
) -> Float["dim1 dim3"]:
pass- 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
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 on it.
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
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+.
Requires Python 3.11+. Runtime type-checking typically requires installing a separate type-checking package.
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"]:
passVerify 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 | permissive license permissive |
| Python support | Supports the current Python release >=3.11 |
| Install friction | Low. Pure-Python wheel |
| Runtime dependencies | 1 packagewadler-lindig |
| Maintenance | Actively maintained 62 days since the last release |
| Last repo commit | |
| First released | |
| Downloads | 7,153,117 / month, #1,779 on PyPI 30-day window, as of 2026-08-14 |
| Known vulnerabilities | None known OSV.dev, checked 2026-08-14 |
| Classifiers | Development Status :: 3 - AlphaIntended Audience :: DevelopersIntended Audience :: Financial and Insurance IndustryIntended Audience :: Information TechnologyIntended Audience :: Science/ResearchLicense :: OSI Approved :: MIT LicenseNatural Language :: EnglishProgramming Language :: Python :: 3Topic :: Scientific/Engineering :: Artificial IntelligenceTopic :: Scientific/Engineering :: Information AnalysisTopic :: Scientific/Engineering :: Mathematics |
Evidence: jaxtyping-0.3.11-py3-none-any.whl
Tags
Let your AI agent find packages like this
Example. Real query, live index.
You found this page by searching. An agent finds it by wishing: SkillFed indexes 14,416 PyPI packages by what they can do, searchable in plain language.
wish › “array shape type annotations”
- jaxtypingProvides type annotations and runtime type-checking for array shape…
- numpydanticNumpydantic adds type annotations and validation for arrays with…
- nptypingProvides type hints and runtime type checking for NumPy arrays and…
Give your agent the search over MCP, or paste the wish link into any chat.
More Artificial Intelligence packages
LiteLLM provides a unified Python interface to call 100+ LLM providers (OpenAI, Anthropic, Gemini, Bedrock, Azure, and others) using OpenAI-compatible API format, available as both a Python SDK and a self-hosted AI Gateway proxy server.
Install it if you need to work with multiple LLM providers or want to centralize LLM routing in your organization.
Client library and CLI tool for downloading, uploading, and managing models, datasets, and repositories on the Hugging Face Hub platform.
Install it if you work with Hugging Face Hub models or datasets.
LangChain provides a framework for building agents and LLM-powered applications by composing language models, tools, and memory through a unified API that abstracts over multiple model providers.
hf-xet provides chunk-based deduplication and efficient file transfer for the Hugging Face Hub, enabling faster uploads and downloads of large files with local disk caching.
Tokenizers converts raw text into token sequences for NLP models, with support for training custom vocabularies and using pre-built tokenizers (BPE, WordPiece) optimized for speed via Rust.
Transformers provides a unified framework for loading, fine-tuning, and running state-of-the-art pretrained models across text, vision, audio, video, and multimodal tasks using PyTorch, JAX, or TensorFlow.
Install it if you need to run or train any transformer-based model for NLP, vision, audio, or multimodal tasks.
See also torchtyping · einshape · chex · equinox · numerary · orbax-export · autoray · nptyping · lineax · numpydantic