jax.lax.pshuffle

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

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

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

参数:

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

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

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

返回:

Array(s) 的形状与 x 相同，其切片沿轴 axis_name 从 x 中根据排列 perm 收集。