jax.tree_util.tree_flatten_with_path

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

別名為 jax.tree.flatten_with_path()。

參數：

• tree (Any)

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

返回類型：

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