mlx.utils.tree_reduce

目录

mlx.utils.tree_reduce#

tree_reduce(fn, tree, initializer=None, is_leaf=None)#

对 Python 树的叶节点应用归约操作。

此函数通过对树的叶节点应用提供的函数 fn，将 Python 树归约到累积结果。

示例

>>> from mlx.utils import tree_reduce
>>> tree = {"a": [1, 2, 3], "b": [4, 5]}
>>> tree_reduce(lambda acc, x: acc + x, tree, 0)
15
参数:

  • fn (callable) – 归约函数，接受两个参数（累加器、当前值）并返回更新后的累加器。

  • tree (Any) – 要进行归约的 Python 树。可以是列表、元组或字典的任意嵌套组合。

  • initializer (Any, optional) – 开始归约的初始值。如果未提供，则使用第一个叶节点值。

  • is_leaf (callable, optional) – 一个函数，用于判断对象是否为叶节点，对叶节点返回 True，否则返回 False。

返回值:

累积值。

返回类型:

Any

目录