Checks that all ``tree`` leaves' shapes have the same prefix. Args: tree: A tree to check. shape_prefix: An expected shape prefix. Raise

Checks that all ``tree`` leaves' shapes have the same suffix. Args: tree: A tree to check. shape_suffix: An expected shape suffix. Raise

Asserts all trees are equal using custom comparator. JIT-friendly.

Checks that trees' leaves have the same shape and dtype. Args: *trees: A sequence of (at least 2) trees to check. Raises: AssertionError

Checks that trees have the same structure and leaves' sizes. Args: *trees: A sequence of (at least 2) trees with array leaves. Raises: A

Checks that the type of all inputs matches specified ``expected_types``. If the expected type is a Python type or abstract dtype (e.g. `np.floating

Raises `ValueError` if `variant_fn` got an unknown argument.

Wraps Chex assert functions to control their common behaviour. Extends the assertion to support the following optional auxiliary kwargs: custom

Clears Chex traces' counter for ``assert_max_traces`` checks. Use it to isolate unit tests that rely on ``assert_max_traces``, by calling it at t

Create a deprecated alias for a function. Example usage: >>> g = create_deprecated_function_alias(f, 'path.f', 'path.g') Args: fun: the de

JAX-friendly wrapper for :py:func:`dataclasses.dataclass`. This wrapper class registers new dataclasses with JAX so that tree utils operate corre

Allows deprecated functions to continue running, with a warning logged.

Disables all Chex assertions. Use wisely.

Enables Chex assertions.

Context manager for patching `jax.jit` and `jax.pmap`. This is a convenience function, equivalent to nested `chex.fake_pmap` and `chex.fake_jit`

Constructs a regexp for the exception message. Args: message: an exception message. Returns: Regexp that ensures the message follows the

Returns a tuple with tracers from a tree.

Wrap chex assertion to only be evaluated if positional args not `None`.

Checks if function is traceable. JAX traces a function when it is wrapped with @jit, @pmap, or @vmap. In other words, this function checks whethe

Converts `jnp.ndarray` to `np.ndarray`.

Resolve a GitHub URL corresponding to Python object.

Returns the number of available device of the given type.

Generates a cartesian product of `params_lists`. See tests from ``variants_test.py`` for examples of usage. Args: *params_lists: A list of p

Disallows JAX compilation for certain backends. Args: allowed: Names of backend platforms (e.g. 'cpu' or 'tpu') for which compilation is

Runs a set of tests for static & value assertions. See `run_test_suite_with_log_abs_fn` for example. Args: make_test_fn: A function th

Forces XLA to use `n` CPU threads as host devices. This allows `jax.pmap` to be tested on a single-CPU platform. This utility only takes effect b