jax.tree.reduce

目录

jax.tree.reduce#

jax.tree.reduce(function: Callable[[T, Any], T], tree: Any, *, is_leaf: Callable[[Any], bool] | None = None) → T[源代码][源代码]#
jax.tree.reduce(function: Callable[[T, Any], T], tree: Any, initializer: T, is_leaf: Callable[[Any], bool] | None = None) → T

对树的叶子调用 reduce()。

参数:

  • function – 归约函数

  • tree – 用于减少的 pytree

  • initializer – 可选的初始值

  • is_leaf – 一个可选指定的函数，将在每次扁平化步骤中被调用。它应返回一个布尔值，指示是否应遍历当前对象，或者是否应立即停止，将整个子树视为叶子。

返回:

减少后的值。

返回类型:

result

示例

>>> import jax
>>> import operator
>>> jax.tree.reduce(operator.add, [1, (2, 3), [4, 5, 6]])
21

参见

目录