jax.tree_util.tree_flatten_with_path

jax.tree_util.tree_flatten_with_path(tree, is_leaf=None, is_leaf_takes_path=False)

别名：jax.tree.flatten_with_path()。

参数:

• tree (Any)

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

• is_leaf_takes_path (bool)

返回类型:

tuple[list[tuple[KeyPath, Any]], PyTreeDef]