torch.linalg.solve

torch.linalg.solve(A, B, *, left=True, out=None) → Tensor

计算具有唯一解的方线性方程组的解。

设 $\mathbb{K}$ 为 $\mathbb{R}$ 或 $\mathbb{C}$，该函数计算 $X \in \mathbb{K}^{n \times k}$ 的解，该解是与 $A \in \mathbb{K}^{n \times n}, B \in \mathbb{K}^{n \times k}$ associated的**线性方程组**，其定义如下：

$$AX = B$$

如果 left= False，则该函数返回矩阵 $X \in \mathbb{K}^{n \times k}$，用于求解以下方程组：

$$XA = B\mathrlap{\qquad A \in \mathbb{K}^{k \times k}, B \in \mathbb{K}^{n \times k}.}$$

当且仅当 $A$ 是 可逆矩阵 时，该线性方程组才有一个解。本函数假设 $A$ 是可逆的。

支持float、double、cfloat和cdouble数据类型的输入。也支持矩阵的批次，如果输入是矩阵的批次，则输出具有相同的批次维度。

设 * 为零个或多个批处理维度，

• 如果 A 的形状为 (*, n, n)，而 B 的形状为 (*, n)（向量批）或形状 (*, n, k)（矩阵批或“多个右侧值”），则此函数返回形状分别为 (*, n) 或 (*, n, k) 的 X。

• 否则，如果 A 的形状为 (*, n, n)，而 B 的形状为 (n,) 或 (n, k)，则 B 将被广播（broadcast）以具有 (*, n) 或 (*, n, k) 的形状。然后，此函数返回由由此产生的线性方程组批次组成的解。

注意

此函数以比单独执行计算更快、更数值稳定的方式计算 X = A.inverse() @ B。

注意

可以通过传递 A 和 B 的转置，并转置此函数返回的输出来计算 $XA = B$ 线性方程组的解。

注意

A 可以是未批处理的 torch.sparse_csr_tensor，但仅限于 left=True。

注意

当输入位于 CUDA 设备上时，此函数会使该设备与 CPU 同步。有关不进行同步的该函数版本，请参阅 torch.linalg.solve_ex()。

另请参阅

torch.linalg.solve_triangular() 计算具有唯一解的三角线性方程组的解。

参数

• A (Tensor) – 形状为 (*, n, n) 的张量，其中 * 是零个或多个批次维度。

• B (Tensor) – 右侧值张量，形状根据上述规则为 (*, n) 或 (*, n, k) 或 (n,) 或 (n, k)

关键字参数

• left (bool, optional) – 是要解方程组 $AX=B$ 还是 $XA = B$。默认为 True。

• out (Tensor, optional) – 输出张量。如果为 None 则忽略。默认为 None。

引发

RuntimeError – 如果 A 矩阵不可逆，或批处理 A 中的任何矩阵不可逆。

示例

>>> A = torch.randn(3, 3)
>>> b = torch.randn(3)
>>> x = torch.linalg.solve(A, b)
>>> torch.allclose(A @ x, b)
True
>>> A = torch.randn(2, 3, 3)
>>> B = torch.randn(2, 3, 4)
>>> X = torch.linalg.solve(A, B)
>>> X.shape
torch.Size([2, 3, 4])
>>> torch.allclose(A @ X, B)
True

>>> A = torch.randn(2, 3, 3)
>>> b = torch.randn(3, 1)
>>> x = torch.linalg.solve(A, b) # b is broadcasted to size (2, 3, 1)
>>> x.shape
torch.Size([2, 3, 1])
>>> torch.allclose(A @ x, b)
True
>>> b = torch.randn(3)
>>> x = torch.linalg.solve(A, b) # b is broadcasted to size (2, 3)
>>> x.shape
torch.Size([2, 3])
>>> Ax = A @ x.unsqueeze(-1)
>>> torch.allclose(Ax, b.unsqueeze(-1).expand_as(Ax))
True