torch.linalg.solve

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

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

令 $\mathbb{K}$ 为实数集 $\mathbb{R}$ 或复数集 $\mathbb{C}$ ，此函数计算与矩阵 $A \in \mathbb{K}^{n \times n}, B \in \mathbb{K}^{n \times k}$ 相关的线性系统的解 $X \in \mathbb{K}^{n \times k}$ ，该系统定义为：

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$ 是可逆的。

支持浮点型、双精度型、复数浮点型和复数双精度型的数据类型。还支持矩阵批处理，如果输入是矩阵批处理，则输出将具有相同的批处理维度。

令 * 代表零个或多个批处理维度。

如果 A 的形状是 (*, n, n)，而 B 的形状为 (*, n)（一批向量）或 (*, n, k)（一批矩阵或多右端），此函数将分别返回 X 形状为 (*, n) 或 (*, n, k)。
否则，如果 A 的形状为 (*, n, n) 且 B 的形状为 (n,) 或 (n, k)，则 B 将被广播到具有形状 (*, 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, 可选) – 是否求解方程组 $AX=B$ 或者 $XA = B$ 。默认值： True。
out (Tensor, optional) – 输出张量。默认为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