mlx.core.distributed.all_sum

all_sum(x: array, *, group: Group | None = None, stream: None | Stream | Device = None) → array

全局归约求和。

对组内所有进程的 x 数组求和。

参数:

• x (array) – 输入数组。

• group (Group) – 参与归约的进程组。如果设置为 None，则使用全局组。默认值: None。

• stream (Stream, optional) – 流或设备。默认为 None，此时使用默认设备的默认流。

返回值:

所有 x 数组的总和。

返回类型:

array