chex
Chex: Testing made fun, in JAX!
Decision gist · record as of 2026-08-14
Yes. Chex is actively maintained, has no known vulnerabilities, and provides essential utilities for JAX development. The low install friction, permissive license, and focus on catching common JAX pitfalls make it a practical addition to any JAX project. Install it if you write JAX code and want better visibility into shape, dtype, and execution-mode issues.AI-flagged interpretation of the facts on this page — verify before relying
Before you install
- Requires Python 3.11 or later; jax and jaxlib must be installed and functional.
- Low friction installation with six runtime dependencies including jax, jaxlib, and numpy.
- Active maintenance with a recent release 63 days ago and ongoing repository activity.
License · maintenance · safety
permissive license (permissive) — Permissive license allows broad use, modification, and distribution with minimal restrictions.
last release 2026-06-12 (63 days) · last repo commit 2026-08-06 · 952 stars
0 known vulnerabilities (OSV.dev, 2026-08-14) · 2,315,756 downloads/mo, #3,144 on PyPI
Alternatives
Verify before relying
pip install chex
import chex
import jax.numpy as jnp
@chex.dataclass
class Parameters:
x: chex.ArrayDevice
y: chex.ArrayDevice
parameters = Parameters(
x=jnp.ones((2, 2)),
y=jnp.ones((1, 2)),
)
chex.assert_equal_shape([parameters.x, parameters.y])- Whether all assertion types work identically in jitted vs non-jitted contexts without additional setup
- Performance overhead of value assertions when using chex.chexify() wrapper in production code
- Compatibility guarantees with specific JAX versions beyond the stated Python 3.11+ requirement
What it is and what it does
Chex is a testing and debugging library for JAX code that bridges the gap between Python's type system and JAX's array constraints. It provides shape, rank, dtype, and device assertions that work within JAX's tracing model, along with utilities to test code across variants (jitted vs non-jitted execution). The library also includes JAX-compatible dataclass implementations and helpers to detect unintended function re-tracing during JIT compilation.
The package is built on top of jax, jaxlib, numpy, and related dependencies. It's designed for machine learning practitioners and researchers who need to catch shape mismatches, type errors, and numerical issues early in development. Assertions can be customized with exception types and messages, and value assertions (those requiring actual tensor values) are supported within jitted functions via the chex.chexify() decorator.
Use it for
- Add shape and dtype validation to JAX functions to catch dimension mismatches before they propagate through a training loop.
- Test the same JAX code path under both jitted and non-jitted execution to verify correctness across compilation modes.
- Detect unintended JIT re-tracing that causes performance degradation by asserting maximum trace counts.
- Validate that model parameters remain finite during training with custom error messages.
- Create JAX-compatible dataclasses that work seamlessly with tree operations.
Worth the install?
AI-flagged interpretation of the facts on this page. Verify before relying on it.
Yes.
Chex is actively maintained, has no known vulnerabilities, and provides essential utilities for JAX development. The low install friction, permissive license, and focus on catching common JAX pitfalls make it a practical addition to any JAX project. Install it if you write JAX code and want better visibility into shape, dtype, and execution-mode issues.
Install
chex on PyPI
Before you install
Low friction installation with six runtime dependencies including jax, jaxlib, and numpy. Active maintenance with a recent release 63 days ago and ongoing repository activity.
Requires Python 3.11 or later; jax and jaxlib must be installed and functional.
License in practice
Permissive license allows broad use, modification, and distribution with minimal restrictions.
Quickstart
pip install chex
import chex
import jax.numpy as jnp
@chex.dataclass
class Parameters:
x: chex.ArrayDevice
y: chex.ArrayDevice
parameters = Parameters(
x=jnp.ones((2, 2)),
y=jnp.ones((1, 2)),
)
chex.assert_equal_shape([parameters.x, parameters.y])Verify before relying
- Whether all assertion types work identically in jitted vs non-jitted contexts without additional setup
- Performance overhead of value assertions when using chex.chexify() wrapper in production code
- Compatibility guarantees with specific JAX versions beyond the stated Python 3.11+ requirement
Package facts
| License | permissive license permissive |
| Python support | Supports the current Python release >=3.11 |
| Install friction | Low. Pure-Python wheel |
| Runtime dependencies | 6 packagesabsl-pytyping_extensionsjaxjaxlibnumpytoolz |
| Maintenance | Actively maintained 63 days since the last release |
| Last repo commit | |
| First released | |
| Downloads | 2,315,756 / month, #3,144 on PyPI 30-day window, as of 2026-08-14 |
| Known vulnerabilities | None known OSV.dev, checked 2026-08-14 |
| Classifiers | Development Status :: 4 - BetaEnvironment :: ConsoleIntended Audience :: DevelopersIntended Audience :: Science/ResearchLicense :: OSI Approved :: Apache Software LicenseOperating System :: OS IndependentProgramming Language :: PythonProgramming Language :: Python :: 3Topic :: Scientific/Engineering :: Artificial IntelligenceTopic :: Software Development :: Libraries :: Python Modules |
Evidence: chex-0.1.92-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 › “jax testing utilities”
- chexChex provides utilities for writing reliable JAX code, including…
- cluCLU provides utilities and abstractions for writing machine learning…
- jraphJraph provides data structures and utilities for building and working…
Give your agent the search over MCP, or paste the wish link into any chat.
More Python Modules packages
Converts domain names between Unicode and ASCII-compatible encoding (Punycode) according to IDNA 2008 and Unicode Technical Standard 46, with security validation and broader script coverage than the standard library.
Install it if you work with internationalized domain names, need to validate domains, or use HTTP clients that depend on it transitively.
Setuptools is a Python build backend and package management tool that handles building, distributing, and installing Python packages, including support for C/C++ extension modules.
PyYAML parses and emits YAML 1.1 data format, enabling serialization and deserialization of configuration files and Python objects to and from human-readable YAML text.
Pydantic validates Python data structures against type hints, coercing and checking input at runtime to ensure it matches a declared schema.
Provides reusable metadata objects for use with PEP-593 `typing.Annotated` to express common constraints like bounds, collection sizes, and predicates on types.
Install it if you use or build libraries that need to express type constraints in a standardized, inspectable way—or if you want to annotate your own types with…
Provides runtime tools to inspect and introspect Python type annotations, enabling programmatic examination of type hints at execution time.
See also distrax · jax · jaxtyping · jaxlib · jax-jumpy · mujoco-mjx · optax · re-assert · torchax · etils