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 上使用。

返回值:

如果 pivoting=False，则返回数组对 (q, r)，否则返回 (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] 的索引向量

Notes

• MAGMA 支持是实验性的 - 有关进一步的假设和限制，请参阅 jax.lax.linalg.eig()。

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