jax.lax.ppermute

jax.lax.ppermute(x, axis_name, perm)

根据排列 perm 执行集体排列。

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

此函数类似于 CollectivePermute HLO。

参数:

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

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

• perm – 一个整数对列表，表示 (source_index, destination_index) 对，它编码了名为 axis_name 的映射轴应该如何混洗。整数值被视为映射轴 axis_name 的索引。任何两个对都不应具有相同的源索引或相同的目标索引。对于轴 axis_name 的每个索引，如果它不对应于 perm 中的目标索引，则结果中的相应值将填充适当类型的零。

返回值:

与 x 形状相同的数组，其沿轴 axis_name 的切片根据排列 perm 从 x 中收集而来。