jax.tree_util.tree_flatten

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

jax.tree.flatten() 的別名。

參數:

• tree (Any)

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

回傳型別:

tuple[list[Leaf], PyTreeDef]