jax.tree
模块#
用于处理树状容器数据结构的实用程序。
The jax.tree
命名空间包含来自 jax.tree_util
的实用程序的别名。
函数列表#
|
在树的叶子上调用 all()。 |
|
展平一个 pytree。 |
|
获取 pytree 的叶子。 |
|
在 pytree 参数上映射一个多输入函数以生成一个新的 pytree。 |
|
在树的叶子上调用 reduce()。 |
|
获取 pytree 的 treedef。 |
|
将具有树结构 (outer, inner) 的树转换为具有结构 (inner, outer) 的树。 |
|
从 treedef 和叶子重建 pytree。 |