jax.lax.linalg.lu_pivots_to_permutation

jax.lax.linalg.lu_pivots_to_permutation(pivots, permutation_size)

将 LU 返回的主元（行交换）转换为置换。

我们构建置换而不是直接将 pivots 应用于矩阵的行，因为 lax 循环不可微分。

参数:

• pivots (ArrayLike) – 形状为 (…, k) 的 int32 数组，表示要执行的行交换

• permutation_size (int) – 输出置换的大小。必须 >= k。

返回:

形状为 (…, permutation_size) 的 int32 数组。

返回类型:

Array