jax.lax.all_to_all

jax.lax.all_to_all(x, axis_name, split_axis, concat_axis, *, axis_index_groups=None, tiled=False)

实现映射轴并映射不同的轴。

如果 x 是一个 pytree，那么结果等效于将此函数映射到树中的每个叶子。

在输出中，输入的映射轴 axis_name 在逻辑轴位置 concat_axis 处实现，并且输入位置 split_axis 处的未映射轴使用名称 axis_name 映射。

映射轴大小的组大小必须等于未映射轴的大小；也就是说，我们必须有 lax.psum(1, axis_name, axis_index_groups=axis_index_groups) == x.shape[axis]。默认情况下，当 axis_index_groups=None 时，这涵盖所有设备。

参数:

• x – 具有名为 axis_name 的映射轴的数组。

• axis_name – 用于命名 pmapped 轴的可哈希 Python 对象（有关更多详细信息，请参阅 jax.pmap() 文档）。

• split_axis – 一个整数，指示要使用名称 axis_name 映射的 x 的未映射轴。

• concat_axis – 一个整数，指示在输出中具体化输入中名为 axis_name 的映射轴的位置。

• axis_index_groups – 可选的列表列表，包含轴索引（例如，对于大小为 4 的轴，[[0, 1], [2, 3]] 将对前两个和后两个副本运行 all_to_all）。组必须恰好覆盖所有轴索引一次，并且所有组的大小必须相同。

• tiled – 当为 True 时，all_to_all 将 split_axis 分成块，并沿着 concat_axis 连接它们。特别是，不会添加或删除任何维度。默认为 False。

返回:

当 tiled 为 False 时，形状由表达式给定的数组

np.insert(np.delete(x.shape, split_axis), concat_axis, axis_size)

其中 axis_size 是输入 x 中名为 axis_name 的映射轴的大小，即 axis_size = lax.psum(1, axis_name)。

否则，返回一个形状类似于输入形状的数组，除了 split_axis 除以轴大小，concat_axis 乘以轴大小。