jax.tree.map#
- jax.tree.map(f, tree, *rest, is_leaf=None)[源代码]#
将一个多输入函数映射到 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) – 一个可选的函数，将在每个展平步骤中调用。它应该返回一个布尔值，指示是否应遍历当前对象，或者是否应立即停止，并将整个子树视为一个叶子节点。
- 返回:
返回一个具有与
tree相同结构的新 pytree，但每个叶子节点的值由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]]