Variant class for typing and string representation.

TestCase generator for chex variants. Supports sharding.

A set of checks imported from checkify.

Thread-safe storage for internal variables used in @chexify.

Counts how often an instance is called.

JAX-friendly wrapper for `dataclasses.dataclass`.

Simple tests for chexify assertions.

Test suite for chexify assertions.

Tests for enable/disable_asserts.

Tests different combinations of base classes for a variants test.

Test chex can be imported correctly.

An enumeration of available Chex variants. Use ``self.variant.type`` to get type of the current test variant. See the docstring of ``chex.variant

Shallow class.

Builder that checks all public symbols are included.

A lightweight utility that maps strings to shape tuples. The most basic usage is: .. code:: >>> dims = chex.Dimensions(B=3, T=5, N=7) # Yo

Nested class.

Checks that Chex raises ValueError when no variants are selected.

Injects a callback into any transformed function. A typical use-case is jax.jit or jax.pmap which is often hidden deep inside the code. This cont

A class for Chex tests that use variants. See the docstring for ``chex.variants`` for more information. Note: ``chex.variants`` returns a genera