jax.tree 模块

用于处理树状容器数据结构的实用程序。

jax.tree 命名空间包含来自 jax.tree_util 的实用程序的别名。

函数列表

all(tree, *[, is_leaf])	在树的叶子上调用 all()。
flatten(tree[, is_leaf])	展平一个 pytree。
flatten_with_path(tree[, is_leaf])	像 tree_flatten 一样扁平化一个 pytree，但也会返回每个叶子的键路径。
leaves(tree[, is_leaf])	获取 pytree 的叶子节点。
leaves_with_path(tree[, is_leaf])	像 tree_leaves 一样获取 pytree 的叶子节点，并返回每个叶子的键路径。
map(f, tree, *rest[, is_leaf])	将一个多输入函数映射到 pytree 参数上，以生成一个新的 pytree。
map_with_path(f, tree, *rest[, is_leaf])	将一个多输入函数映射到 pytree 的键路径和参数上，以生成一个新的 pytree。
reduce()	在树的叶子上调用 reduce()。
structure(tree[, is_leaf])	获取 pytree 的 treedef。
transpose(outer_treedef, inner_treedef, ...)	将具有树结构 (outer, inner) 的树转换为具有结构 (inner, outer) 的树。
unflatten(treedef, leaves)	从 treedef 和叶子重建一个 pytree。