jax.tree.leaves

内容

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

获取 pytree 的叶子节点。

参数：

tree (Any) – 要获取叶子节点的 pytree
is_leaf (Callable[[Any], bool] | None | None) – 可选指定的函数，将在每个扁平化步骤中调用。它应该返回一个布尔值，指示是否应该遍历当前对象，或者是否应该立即停止，并将整个子树视为叶子节点。

返回值：

树叶节点列表。

返回类型：

leaves

示例

>>> import jax
>>> jax.tree.leaves([1, (2, 3), [4, 5]])
[1, 2, 3, 4, 5]

另请参阅