Pruning API Reference

Traverse a PyTree structure and prune out leaves that do not satisfy a predicate.

Comparison

• JAX: Similar to jax.tree_util.tree_map, but dynamically removes elements from the PyTree structure instead of keeping a 1:1 mapping.
• Equinox: Allows selective filtering of components, but unlike Equinox's filtering which replaces leaves with None or leaves the tree structure unchanged, tree_map_prune completely removes (collapses) elements from lists, tuples, and dictionaries.

Example

import jaxtree_extensions as jte

tree = {"a": [1, 2, 3], "b": 4}
# Map x * 10, keeping only values that are not divisible by 20
pruned, treedef = jte.tree_map_prune(
    lambda x: x * 10,
    tree,
    keep_fn=lambda x: x % 20 != 0
)

# pruned = {'a': [10, 30], 'b': None}

# Reconstruct the original structure
original = treedef.unflatten(pruned, fill_value=0)
# original = {'a': [10, 0, 30], 'b': 0}

jaxtree_extensions.tree_map_prune(f: collections.abc.Callable[..., typing.Any], tree: typing.Any, *rest: typing.Any, keep_fn: collections.abc.Callable[[typing.Any], bool], is_leaf: collections.abc.Callable[[typing.Any], bool] | None = None) -> tuple[typing.Any, jaxtree_extensions.prune.PrunedTreeDef]

Map a function over a PyTree, pruning leaves that don't match a predicate.

For collapsable PyTree containers (dicts, lists, tuples), pruned leaves are completely filtered out. For non-collapsable or custom PyTree containers, pruned leaves are replaced with None.

Parameters:

Name	Type	Description	Default
f	Callable[..., Any]	The function to map over PyTree leaves.	required
tree	Any	The primary PyTree to traverse.	required
*rest	Any	Additional matching PyTrees to map over in parallel.	required
keep_fn	Callable[[Any], bool]	A predicate called on each mapped leaf to determine if it should be kept. Returns True to keep, False to prune.	required
is_leaf	Callable[[Any], bool] | None	A predicate called on each node to determine if it should be treated as a leaf instead of traversed.	None

Returns:

Type	Description
tuple[Any, PrunedTreeDef]	A tuple containing: pruned_tree: The mapped and pruned PyTree. treedef: A PrunedTreeDef structure definition that can reconstruct the original tree structure (e.g. replacing pruned nodes with custom values).

jaxtree_extensions.PrunedTreeDef

Structure definition for reconstructing pruned JAX PyTrees.

Allows restoring the original PyTree structure from a pruned tree, filling pruned leaves with a specific value or a matching template PyTree.

Parameters:

Name	Type	Description	Default
template	Any	The template PyTree representing the original structure.	required
is_leaf_fn	Callable[[Any], bool] | None	A predicate called on each node to determine if it should be treated as a leaf.	None

unflatten(pruned_tree: typing.Any, fill_value: typing.Any = None) -> typing.Any

Reconstruct the original PyTree structure from a pruned tree.

Parameters:

Name	Type	Description	Default
pruned_tree	Any	The pruned PyTree returned by tree_map_prune.	required
fill_value	Any	The value to fill pruned leaf positions with. Can be a single value (like None), or a PyTree matching the original structure.	None

Returns:

Type	Description
Any	The reconstructed PyTree with the original structure.