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 | None) – 可选指定的函数，将在每个展平步骤中调用。它应该返回一个布尔值，指示展平是否应遍历当前对象，或者是否应立即停止，并将整个子树视为叶子。
- 返回:
一个新的 pytree，其结构与
tree相同，但每个叶子上的值由f(x, *xs)给出，其中x是tree中相应叶子上的值，xs是rest中相应节点上的值的元组。- 返回类型:
Any
示例
>>> 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]]