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 形状相同的数组(s)，沿着轴 axis_name 从 x 中收集切片，根据排列 perm。