PyTorch torch.linalg.lstsq Function


Pytorch torch 参考手册Pytorch torch Reference Manual

torch.linalg.lstsqThis is a function in PyTorch's linear algebra module used to solve linear least squares problems. It solves argmin_x ||AX - B||_F.

Function Definition

torch.linalg.lstsq(A, B, rcond=None, driver=None)

Parameters:

  • A(Tensor): Coefficient matrix.
  • B(Tensor): Right-hand side matrix or vector.
  • rcond(float, optional): Condition number used for truncating singular values.
  • driver(str, optional): Solver selection; choices are 'gels', 'gelsd', 'gelsy'.

Return Value:

  • torch.Tensor: Returns the least squares solution.

Usage Example

Example

import torch

# Create coefficient matrix and right-hand side vector
A = torch.tensor([[1.0, 1.0], [1.0, 2.0], [1.0, 3.0]], dtype=torch.float64)
B = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64)

# Solve least squares
X = torch.linalg.lstsq(A, B).solution

print("Coefficient matrix A:")
print(A)
print("nRight-hand side vector B:")
print(B)
print("nLeast squares solution X:")
print(X)

The output is:

系数矩阵 A:
tensor([[1., 1.],
        [1., 2.],
        [1., 3.]], dtype=torch.float64)
右侧向量 B:
tensor([1., 2., 3.], dtype=float64)
最小二乘解 X:
tensor([0.0000, 1.0000], dtype=float64)

Pytorch torch 参考手册Pytorch torch Reference Manual

Other Extensions