Raises a RuntimeError if not overriden or redefined.

Initializes test data.

Wrapper that forces axis_index_groups to be None. This is to avoid problems within fake_pmap where parallel operations are performed with vmap, r

Asserts that a function can be jitted or not. Args: fn: The function to be tested fn_input: Input to pass to the function is_jitted: As

Checks that all trees are equal as per the custom comparator for leaves. Args: equality_comparator: A custom function that accepts two leaves a

Asserts whether a function can be pmapped or not. Args: fn: The function to be tested fn_input: Input to pass to the function is_pmappe

Returns a generator with standard tests. For internal usage. Allows to dynamically generate common tests. See tests' names and comments for more

Asserts that `pass_fn(tensor.shape[axis])` passes. Used to implement ==, >, >=, <, <= checks. Args: tensor: A JAX array. axis: An intege

Generates valid and invalid inputs for log_abs_fn and runs the tests.

Checks that assert_fn correctly processes invalid args' structs.

Checks that ``x`` is a scalar, as defined in `pytypes.py` (int or float). Args: x: An object to check. Raises: AssertionError: If ``x``

Checks the error and converts it to chex format.

Serialization-friendly version of dataclasses.fields for instances.

Returns True if `actual_shape` is compatible with `expected_shape`.

Checks that all `tree`'s leaves are n-dimensional arrays (tensors). Args: tree: A tree to assert. Raises: AssertionError: If the tree co

Exposes dataclass as ``collections.abc.Mapping`` descendent. Allows to traverse dataclasses in methods from `dm-tree` library. NOTE: changes dat

Register an existing dataclass so JAX knows how to handle it. This means that functions in jax.tree_util operate over the fields of the dataclass

Returns the flat size of a given named shape, i.e. prod(shape).

Patch the given backend restriction onto getattr(compiler, name).

Returns True if `actual_shape` is compatible with `expected_shape`.

Implements `variants` and `all_variants`.

Appends a future annotations import to the file at the given path.

Number of representable bf16 points between each item in x and y.

Generates shallow and nested mappable dataclasses.

Initializes common test cases.

Generate chex variants for a single test.

Constructs a host assertion given `assert_fn`. This wrapper should only be applied to the assertions that are either a) never used in jitted co

Set a name for the generated test.

Merge variants.

Checks that a scalar is non-negative. Args: x: A value to check. Raises: AssertionError: If ``x`` is not a scalar or negative.

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

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

Checks that trees have the same structure. Args: *trees: A sequence of (at least 2) trees to assert equal structure between. Raises: Val

Waits until all asynchronous checks complete. See `chexify` for more detail.

Wraps a transformed function `fn` to enable Chex value assertions. Chex value/runtime assertions access concrete values of tensors (e.g. `assert_

Converts a path from jax.tree_util to one from dm-tree.

Converts varargs+kwargs function arguments into varargs only.

Counts the number of times the function was called.

Context manager for patching `jax.jit` with the identity function. This is intended to be used as a debugging tool to programmatically enable or

Context manager for patching `jax.pmap` with `jax.vmap`. This is intended to be used as a debugging tool to programmatically replace pmap transfo

Returns `(name, module)` for all submodules under `root_module`.

Returns a test class derived from the specified bases.

Constructs an error message for the chexify exception.

Parses number of CPUs from the XLA environment flags.

Returns the latest non-chex frame from the call stack.

Checks whether a tree contains any tracers.

Forwards class to dataclasses's wrapper and registers it with JAX.