Install & Compatibility
Where this runs
tested against v0.1.90 · pip install
no network on importno background threads
Install × environment matrix
Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux
muslpy 3.10–3.95 runs
build_error
glibcpy 3.10–3.95 runs
installs and imports cleanly · install 12.7s · import 2.686s · 565MB
590MB installed
● package 590MB
Code
Verified usage
Verified import paths — ran on the pinned version, not inferred.
dataclass
✓ from chex import dataclass
✗ from dataclasses import dataclass
Chex provides a JAX-compatible dataclass implementation.
ArrayDevice
✓ from chex import ArrayDevice
Type hint for JAX array on a device.
assert_tree_all_finite
✓ from chex import assert_tree_all_finite
chexify
✓ from chex import chexify
variants
✓ from chex import variants
assert_max_traces
✓ from chex import assert_max_traces
This quickstart demonstrates defining a JAX-compatible dataclass using `chex.dataclass`, performing a JAX `tree_map` operation on it, and using `chex.assert_tree_all_finite` within a JIT-compiled function by decorating it with `chex.chexify`.
import chex
import jax
import jax.numpy as jnp
# Define a JAX-friendly dataclass
@chex.dataclass
class Parameters:
x: chex.ArrayDevice
y: chex.ArrayDevice
# Create an instance
params = Parameters(x=jnp.ones((2, 2)), y=jnp.ones((1, 2)))
# Dataclasses can be treated as JAX pytrees
transformed_params = jax.tree_util.tree_map(lambda val: 2.0 * val, params)
print(f"Original params: {params.x}\nTransformed params: {transformed_params.x}")
# Use an assertion
def my_func(val):
chex.assert_tree_all_finite(val)
return val * 2
# Assertions can be used within jitted functions with chexify
@chex.chexify
@jax.jit
def jitted_func(val):
return my_func(val)
# This will pass
jitted_func(jnp.array([1.0, 2.0]))
# This would fail (if uncommented) because of NaN values
# try:
# jitted_func(jnp.array([1.0, jnp.nan]))
# except chex.errors.ChexTypeError as e:
# print(f"Caught expected error: {e}")
Debug
Known issues
breakingChex's `mappable_dataclass` and `dataclass` implementations do not support positional arguments for construction, unlike standard Python dataclasses. Arguments must be provided as keyword arguments, similar to a dictionary constructor.fixAlways initialize `chex.dataclass` instances using keyword arguments (e.g., `MyParams(x=1, y=2)` instead of `MyParams(1, 2)`).
affects: All versions
breakingChex has transitioned from relying on `dm-tree` to using JAX's native `jax.tree_util` for PyTree operations. As a result, `None` values are no longer treated as distinct leaves by `chex` tree assertions by default.fixTo explicitly check for `None`s in PyTrees, use `chex.assert_tree_no_nones()` or similar specific assertions.
affects: Versions migrating from `dm-tree` usage (check release notes for specific version, but generally recent versions).
gotchaWhen using `chex.chexify()` with JIT-compiled functions, assertions might run asynchronously. This means errors may not be raised immediately but potentially at a later line or function call. For reliable testing, especially when expecting an assertion to fail, you might need to explicitly wait for checks to complete.fixAfter calling a `chexify`'d function, call `.wait_checks()` on the function object (e.g., `jitted_func.wait_checks()`) to ensure all asynchronous assertions have completed and raised any errors.
affects: All versions using `chex.chexify()` for async assertions.
gotchaThe `chex.assert_max_traces()` decorator (and similar tracing assertions like `assert_max_retraces`) expects to wrap a pure Python function, not an already JIT-compiled function. Applying it to a function that has already been decorated with `jax.jit` will likely lead to incorrect behavior or assertion failures.fixEnsure `chex.assert_max_traces()` is applied *before* `jax.jit` if both are used, or wrap a non-jitted function directly. For example: `@jax.jit @chex.assert_max_traces(n=1) def fn(...)`.
affects: All versions.
Errors
Common errors & fixes
ModuleNotFoundError: No module named 'chex'
The 'chex' library is not installed in your current Python environment or Python cannot find it in the configured paths.
fixInstall the library using pip: `pip install chex`. If using a virtual environment, ensure it is activated before installation.
AssertionError: expected shapes should be a list or tuple of ints, got Ellipsis
When using `chex.assert_shape` to specify a wildcard dimension with Ellipsis (`...`), it must be wrapped within a tuple, even if it's the only element, e.g., `(...,)`.
fixModify the `assert_shape` call to wrap the Ellipsis in a tuple: `chex.assert_shape(x, (...,))` instead of `chex.assert_shape(x, ...)`.
AttributeError: module 'chex' has no attribute 'warn_deprecated_function'
This typically occurs when a dependent library (like `optax` or `scvi-tools`) requires a newer version of `chex` that includes the `warn_deprecated_function`, but an older version of `chex` is currently installed.
fixUpgrade the `chex` library to a compatible version, often by running `pip install --upgrade chex` or `pip install chex>=0.1.86` to meet the dependency requirements of the calling library.
[Chex] Assertion assert_equal_shape failed: Arrays have different shapes: [(X, Y), (Z,)]
The JAX arrays being compared by `chex.assert_equal_shape` or `chex.assert_shape` have different dimensions or sizes, which violates the assertion.
fixExamine the shapes of the arrays involved in the assertion to ensure they match the expected dimensions. Adjust your code to correctly reshape the arrays or modify the expected shape argument passed to `chex.assert_equal_shape`/`chex.assert_shape`.
Upgrade
Version history
0.1.92latest on PyPI · released Jun 12, 2026
Audit
Dependencies
absl-pyrequiredRuntime dependency
typing_extensionsrequiredRuntime dependency for type hints
jaxrequiredCore JAX dependency
jaxlibrequiredCore JAX dependency
numpyrequiredNumerical operations
toolzrequiredFunctional utilities