jax.tree.flatten#

jax.tree.flatten(tree, is_leaf=None)[源代码]#

展平一个 pytree。

展平顺序（即输出列表中元素的顺序）是确定的，对应于从左到右的深度优先树遍历。

参数:

  • tree (Any) – 要展平的 pytree。

  • is_leaf (Callable[[Any], bool] | None | None) – 一个可选指定函数，该函数将在每个展平步骤中被调用。它应该返回一个布尔值，true 表示停止遍历，并将整个子树视为叶子，false 表示展平应该遍历当前对象。

返回:

一个对，其中第一个元素是叶子值的列表，第二个元素是表示展平树结构的 treedef。

返回类型:

元组[列表[tree_util.Leaf], tree_util.PyTreeDef]

示例

>>> import jax
>>> vals, treedef = jax.tree.flatten([1, (2, 3), [4, 5]])
>>> vals
[1, 2, 3, 4, 5]
>>> treedef
PyTreeDef([*, (*, *), [*, *]])

另请参阅