jax.tree.map#
- jax.tree.map(f, tree, *rest, is_leaf=None)[source]#
将多输入函数映射到 pytree 参数上以生成新的 pytree。
- 参数:
f (Callable[..., Any]) – 一个接受
1 + len(rest)个参数的函数，将应用于 pytree 的相应叶子。tree (Any) – 一个要进行映射的 pytree，每个叶子提供给
f的第一个位置参数。rest (Any) – 一个 pytree 元组，其中每个 pytree 都具有与
tree相同的结构，或以tree作为前缀。is_leaf (Callable[[Any], bool] | None) – 一个可选指定函数，将在每个展平步骤中调用。它应返回一个布尔值，指示展平是否应遍历当前对象，或者是否应立即停止，并将整个子树视为叶子。
- 返回:
一个新的 pytree，其结构与
tree相同，但每个叶子的值由f(x, *xs)给出，其中x是tree中相应叶子的值，xs是rest中相应节点的值元组。- 返回类型:
任意类型
示例
>>> import jax >>> jax.tree.map(lambda x: x + 1, {"x": 7, "y": 42}) {'x': 8, 'y': 43}
如果传入多个输入，树的结构将取自第一个输入；后续输入只需以
tree为前缀即可。>>> jax.tree.map(lambda x, y: [x] + y, [5, 6], [[7, 9], [1, 2]]) [[5, 7, 9], [6, 1, 2]]