Map Enumerate API Reference

Map a function over a PyTree, passing a sequential flat index to each leaf.

Comparison

• JAX: Identical to jax.tree_util.tree_map, but provides a running flat leaf index starting from 0 to the mapping function f(index, leaf), saving the boilerplate of manually tracking state in a mutable list or closure.

Example

import jaxtree_extensions as jte

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

# Tag each leaf with its flat index
res = jte.tree_map_enumerate(lambda idx, x: (idx, x), tree)
# res = {"a": [(0, 10), (1, 20)], "b": (2, 30)}

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

Map a function over leaves, passing the flat leaf index as the first argument.

Parameters:

Name	Type	Description	Default
f	Callable[..., Any]	The function to map, called as f(idx, leaf, *rest_leaves).	required
tree	Any	The primary PyTree.	required
*rest	Any	Additional matching PyTrees.	required
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 mapped PyTree.