jax.lax.linalg.qr

jax.lax.linalg.qr(x: ArrayLike, *, pivoting: Literal[False], full_matrices: bool = True, use_magma: bool | None = None) → tuple[Array, Array]

jax.lax.linalg.qr(x: ArrayLike, *, pivoting: Literal[True], full_matrices: bool = True, use_magma: bool | None = None) → tuple[Array, Array, Array]

jax.lax.linalg.qr(x: ArrayLike, *, pivoting: bool = False, full_matrices: bool = True, use_magma: bool | None = None) → tuple[Array, Array] | tuple[Array, Array, Array]

QR 分解。

计算 QR 分解

\[A = Q \, R\]

对于矩阵 \(A\)，其中 \(Q\) 是一个酉（正交）矩阵，而 \(R\) 是一个上三角矩阵。

参数:

• x – 形状为 [..., m, n] 的一批矩阵。

• pivoting – 允许 QR 分解揭示秩。 如果 True，则计算列主元分解 A[:, P] = Q @ R，其中选择 P 使得 R 的对角线是非递增的。 目前仅在 CPU 和 GPU 后端支持。

• full_matrices – 确定是返回完整矩阵还是缩减矩阵；请参见下文。

• use_magma – 在本地覆盖 jax_use_magma 标志。 如果 True，则使用 MAGMA 计算主元 qr 分解。 如果 False，则使用主机 CPU 上的 LAPACK 进行计算。 如果 None（默认），则行为由 jax_use_magma 标志控制。 此参数仅在 GPU 上使用。

返回:

一对数组 (q, r)，如果 pivoting=False，否则为 (q, r, p)。

数组 q 是一个酉（正交）矩阵，如果 full_matrices=True，则形状为 [..., m, m]，如果 full_matrices=False，则形状为 [..., m, min(m, n)]。

数组 r 是一个上三角矩阵，如果 full_matrices=True，则形状为 [..., m, n]，如果 full_matrices=False，则形状为 [..., min(m, n), n]。

数组 p 是形状为 […, n] 的索引向量

注意事项

• MAGMA 支持是实验性的 - 请参阅 jax.lax.linalg.eig() 了解更多假设和限制。

• 如果 jax_use_magma 设置为 "auto"，如果可以找到该库，并且输入矩阵足够大（至少有 2048 列），则将使用 MAGMA 实现。