mlx.core.distributed.all_gather

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

从所有进程收集数组。

从组内所有进程收集 x 数组，并将它们沿第一个轴连接起来。所有数组应具有相同的形状。

参数:

• x (array) – 输入数组。

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

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

返回:

所有 x 数组的连接结果。

返回类型:

array