Creates a new OnCallOfTransformedFunction context manager. Args: fn_transformation: identifier of the function transformation e.g.

Generate chex variants for each test case.

No-op.

A jittable version of `_assert_tree_all_finite_static`.

Checks that all leaves in a tree are finite. Args: tree_like: A pytree with array leaves. Raises: AssertionError: If any leaf in ``tree_

A jittable version of `_assert_trees_all_close_static`.

Checks that all trees have leaves with approximately equal values. This compares the difference between values of actual and desired up to ``ato

A dummy jittable version of `_assert_trees_all_close_ulp_static`. JAX does not yet have a native version of assert_array_max_ulp, so at the momen

Checks that tree leaves differ by at most `maxulp` Units in the Last Place. This is the Chex version of np.testing.assert_array_max_ulp. Asserti

A jittable version of `_assert_trees_all_equal_static`.

Checks that all trees have leaves with *exactly* equal values. If you are comparing floating point numbers, an exact equality check may not be ap

Creates a chex dataclass from a flatten jax.tree_util representation.

Fake implementation of pmap using vmap.

Returns a transformed version of the given function.

Returns result of the returned function and calls the callback.

Variant that applies `jax.device_put` to the args of fn.

Variant that applies `jax.jit` to fn.

Variant that applies `jax.pmap` to fn. Args: fn: A function to wrap. broadcast_args_to_devices: Whether to broadcast `fn` args to pmap form

Variant that applies `jax.device_get` to the args of fn.

Variant that does not apply `jax.jit` to a fn (identity).

Equivalent to ``chex.variants`` but with flipped defaults.

Checks that ``tensor.shape[axis] == expected``. Args: tensor: A JAX array. axis: An integer specifying which axis to assert. expected:

Checks that ``tensor.shape[axis] > val``. Args: tensor: A JAX array. axis: An integer specifying which axis to assert. val: A value ``t

Checks that ``tensor.shape[axis] >= val``. Args: tensor: A JAX array. axis: An integer specifying which axis to assert. val: A value ``

Checks that ``tensor.shape[axis] < val``. Args: tensor: A JAX Array. axis: An integer specifiying which axis to assert. val: A value ``

Checks that ``tensor.shape[axis] <= val``. Args: tensor: A JAX array. axis: An integer specifying which axis to assert. val: A value ``

Checks if ``inputs`` is a collection of arrays.

Checks that `n` devices of a given type are available. Args: n: A required number of devices of the given type. devtype: A type of devices,

Checks that the two objects are equal as determined by the `==` operator. Arrays with more than one element cannot be compared. Use ``assert_tree

Checks that all arrays have the same rank. Args: inputs: A collection of arrays. Raises: AssertionError: If the ranks of all arrays do n

Checks that all arrays have the same shape. Args: inputs: A collection of arrays. dims: An optional integer or sequence of integers. If not

Checks that the leading ``prefix_dims`` dims of all inputs have same shape. Args: inputs: A collection of input arrays. prefix_len: A numbe

Checks that the final ``suffix_len`` dims of all inputs have same shape. Args: inputs: A collection of input arrays. suffix_len: A number o

Checks that all arrays have the same size. Args: inputs: A collection of arrays. Raises: AssertionError: If the size of all arrays do no

Checks that one and only one of the arguments is `None`. Args: first: A first object. second: A second object. Raises: AssertionErro

Checks that at least one GPU device is available. Args: backend: A type of backend to use (uses JAX default if not provided). Raises: As

Checks that an array of ``shape_a`` is broadcastable to one of ``shape_b``. Args: shape_a: A shape of the array to check. shape_b: A target

Checks that ``numerator`` is divisible by ``denominator``. Args: numerator: A numerator. denominator: A denominator. Raises: Asserti

Asserts all leaves are equal using custom comparator. Not jittable.

Checks that a function is traced at most `n` times (inclusively). JAX re-traces jitted functions every time the structure of passed arguments cha

Checks that at least one of the arguments is not `None`. Args: first: A first object. second: A second object. Raises: AssertionErro

Checks that autodiff and numerical gradients of a function match. Args: f: A function to check. f_args: Arguments of the function. orde

Checks that the rank of all inputs matches specified ``expected_ranks``. Valid usages include: .. code-block:: python assert_rank(x, 0)

Checks that argument is a scalar within segment (by default). Args: x: An object to check. min_: A left border of the segment. max_: A

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

Checks that a scalar is positive. Args: x: A value to check. Raises: AssertionError: If ``x`` is not a scalar or strictly positive.

Checks that the shape of all inputs matches specified ``expected_shapes``. Valid usages include: .. code-block:: python assert_shape(x, ())

Checks that the size of all inputs matches specified ``expected_sizes``. Valid usages include: .. code-block:: python assert_size(x, 1)

Checks that at least one TPU device is available. Args: backend: A type of backend to use (uses JAX default if not provided). Raises: As

Checks that all leaves are ndarrays residing in device memory (in HBM). Sharded DeviceArrays are disallowed. Args: tree: A tree to assert.

Checks that all leaves are ndarrays residing in the host memory (on CPU). This assertion only accepts trees consisting of ndarrays. Args: tr

Checks that all leaves are ndarrays sharded across the specified devices. Args: tree: A tree to assert. devices: A list of devices which th

Checks that a tree does not contain `None`. Args: tree: A tree to assert. Raises: AssertionError: If the tree contains at least one `Non

Checks that all ``tree`` leaves match the ``expected_shape``. Args: tree: A tree to check. expected_shape: An expected shape. See ``chex.as