jax.tree_util.tree_leaves

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

jax.tree.leaves()的别名。

参数：

• tree (Any)

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

返回类型：

list[Leaf]