chex

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

Chex provides utilities for writing reliable JAX code, including assertions for tensor properties, debugging helpers, and test variants to validate code across different JAX execution modes.

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

pip install chex

uv add chex

poetry add chex

Installing chex

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.

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

Requires Python 3.11 or later; jax and jaxlib must be installed and functional.

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

Evidence: chex-0.1.92-py3-none-any.whl

Keywords: jax, testing, debugging, python, machine learning

Tags

More Python Modules packages