jax.lax.pswapaxes

jax.lax.pswapaxes(x, axis_name, axis, *, axis_index_groups=None)

将 pmapped 轴 axis_name 与未映射轴 axis 交换。

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

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

此函数是 all_to_all 的一个特例，其中输入的 pmapped 轴放置在输出中的 axis 位置。也就是说，它等价于 all_to_all(x, axis_name, axis, axis)。

参数:

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

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

• axis – int，指示 x 中要使用名称 axis_name 映射的未映射轴。

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

返回值:

与 x 具有相同形状的数组。