jax.lax.pshuffle

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

jax.lax.ppermute 的便捷包装器，使用备用排列编码

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

参数：

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

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

• perm – 整数列表，对名为 axis_name 的轴要应用的排列进行编码，以便输出在轴索引 i 处来自输入在轴索引 perm[i] 处。对于轴大小 N，[0, N) 中的每个整数都应恰好包含一次。

返回值：

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