Path Reduction API Reference

Reduce a PyTree to a single value, passing the key path of each leaf to the reduction function.

Comparison

• JAX: Extends jax.tree_util.tree_reduce by using path-aware traversal. Similar to how jax.tree_util.tree_map_with_path extends jax.tree_util.tree_map, tree_reduce_with_path allows tracking leaf context paths (e.g. key-path tuples) during the aggregation phase.

Example

import jaxtree_extensions as jte

tree = {"a": [10, 20], "b": 30}

# Collect string representations of all leaf paths
paths = jte.tree_reduce_with_path(
    lambda acc, path, leaf: acc + [jte.path_to_str(path)],
    tree,
    initializer=[]
)
# paths = ["/a/0", "/a/1", "/b"]

jaxtree_extensions.tree_reduce_with_path(f: collections.abc.Callable[[typing.Any, tuple[typing.Any, ...], typing.Any], typing.Any], tree: typing.Any, initializer: typing.Any = Ellipsis, is_leaf: collections.abc.Callable[[typing.Any], bool] | None = None) -> typing.Any

Reduce a PyTree over its leaves, passing their key paths to the reduction function.

Parameters:

Name	Type	Description	Default
f	Callable[[Any, tuple[Any, ...], Any], Any]	The reduction function, called as f(accumulator, path, leaf).	required
tree	Any	The PyTree to reduce.	required
initializer	Any	The initial accumulator value. If not specified, the first leaf value is used as the initial accumulator.	Ellipsis
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
Any	The reduced accumulator value.

Raises:

Type	Description
ValueError	If the PyTree is empty and no initializer is provided.