jax.tree_util.tree_leaves

目录

jax.tree_util.tree_leaves#

jax.tree_util.tree_leaves(tree, is_leaf=None)[source]#

jax.tree.leaves() 的别名。

参数:
  • tree (Any)

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

返回类型:

list[Leaf]