jax.tree_util.tree_flatten_with_path# jax.tree_util.tree_flatten_with_path(tree, is_leaf=None)[source]# Alias of jax.tree.flatten_with_path(). Parameters: tree (Any) is_leaf (Callable[[Any], bool] | None | None) Return type: tuple[list[tuple[KeyPath, Any]], PyTreeDef]
