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

另請參閱