Comparison API Reference

Compare two PyTrees structure, types, and array closeness, generating a detailed difference report.

Comparison

• Chex: Similar to chex.assert_trees_all_close or other validation libraries, but instead of raising an immediate assertion error on the first mismatch, tree_compare performs a complete traversal and returns a structured TreeDiff diagnostic object describing all mismatches in detail.

Example

import jaxtree_extensions as jte
import jax.numpy as jnp

tree_a = {"a": [1, 2], "b": jnp.ones((2, 2))}
tree_b = {"a": [1.0, 2], "b": jnp.ones((2, 2)) + 0.1}

diff = jte.tree_compare(tree_a, tree_b, rtol=1e-3, atol=1e-3)
if diff.has_changes:
    print(diff.report())
# Output shows the type mismatch at 'a/0' and the value difference at 'b'

jaxtree_extensions.tree_compare(tree_a: Any, tree_b: Any, rtol: float = 1e-05, atol: float = 1e-08, equal_nan: bool = True, is_leaf: Callable[[Any], bool] | None = None) -> TreeDiff

Compare two PyTrees and return a detailed diagnostic diff report.

Checks structure, node types, and leaf values (using array closeness for JAX and NumPy arrays, and strict equality for other leaf types).

Parameters:

Name	Type	Description	Default
tree_a	Any	The first PyTree.	required
tree_b	Any	The second PyTree.	required
rtol	float	Relative tolerance for numerical closeness.	1e-05
atol	float	Absolute tolerance for numerical closeness.	1e-08
equal_nan	bool	Whether to treat NaNs as equal.	True
is_leaf	Callable[[Any], bool] | None	A predicate called on each node to determine if it should be treated as a leaf.	None

Returns:

Type	Description
TreeDiff	A TreeDiff object representing the diagnostic report.

jaxtree_extensions.TreeDiff

Diagnostic report of differences between two PyTrees.

Parameters:

Name	Type	Description	Default
mismatches	list[Mismatch]	A list of recorded Mismatch objects.	required

has_changes property

Return True if any differences were found.

report() -> str

Return a formatted, readable string of all differences.

Returns:

Type	Description
str	A multi-line diagnostic report.