optax.tree_utils.tree_allclose

Contents

optax.tree_utils.tree_allclose#

optax.tree_utils.tree_allclose(
a: Any,
b: Any,
rtol: jax.typing.ArrayLike = 1e-05,
atol: jax.typing.ArrayLike = 1e-08,
equal_nan: bool = False,
)[source]#

Check whether two trees are element-wise approximately equal within a tolerance.

See jax.numpy.allclose() for the equivalent on arrays.

Parameters:

  • a – a tree

  • b – a tree

  • rtol – relative tolerance used for approximate equality

  • atol – absolute tolerance used for approximate equality

  • equal_nan – boolean indicating whether NaNs are treated as equal

Returns:

a boolean value.

Contents