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 | None) – 一个可选的函数，将在每次扁平化步骤中被调用。它应该返回一个布尔值，指示是否应该遍历当前对象，或者是否应该立即停止，并将整个子树视为一个叶子节点。

返回:

一个与 tree 具有相同结构的新 pytree，但每个叶子节点的值由 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]]

另请参阅

• jax.tree.leaves()

• jax.tree.reduce()