jax.tree 模块

内容

jax.tree 模块#

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

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

函数列表#

all(tree, *[, is_leaf])

在树的叶子上调用 all()。

flatten(tree[, is_leaf])

展平一个 pytree。

leaves(tree[, is_leaf])

获取 pytree 的叶子。

map(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。