jax.tree_util.tree_leaves_with_path

jax.tree_util.tree_leaves_with_path(tree, is_leaf=None)

Alias of jax.tree.leaves_with_path().

Parameters:

• tree (Any)

• is_leaf (Callable[[Any], bool] | None | None)

Return type:

list[tuple[KeyPath, Any]]