PyTorch – torch.linalg.solve() Method

To solve a square system of linear equations with unique solution, we could apply the torch.linalg.solve() method. This method takes two parameters −
first, the coefficient matrix A, and
first, the coefficient matrix A, and
second, the right-hand tensor b.
second, the right-hand tensor b.
Where A is a square matrix and b is a vector. The solution is unique if A invertible. We can solve a number of systems of linear equations. In this case, A is a batch of square matrices and b is a batch of vectors.
torch.linalg.solve(A, b)
A – Square matrix or batch of square matrices. It is the coefficient matrix of system of linear equations.
A – Square matrix or batch of square matrices. It is the coefficient matrix of system of linear equations.
b – Vector or a batch of vectors. It's the right-hand tensor of the linear system.
b – Vector or a batch of vectors. It's the right-hand tensor of the linear system.
It returns a tensor of the solution of the system of linear equations.
Note − This method assumes that the coefficient matrix A is invertible. If it is not invertible, a Runtime Error will be raised.
We could use the following steps to solve a square system of linear equations.
Import the required library. In all the following examples, the required Python library is torch. Make sure you have already installed it.
Import the required library. In all the following examples, the required Python library is torch. Make sure you have already installed it.
import torch
Define a Coefficient matrix and the right-hand side tensor for the given square system of linear equations.
Define a Coefficient matrix and the right-hand side tensor for the given square system of linear equations.
A = torch.tensor([[2., 3.],[1., -2.]])
b = torch.tensor([3., 0.])
Compute the unique solution using torch.linalg.solve(A,b).
Coefficient matrix A must be invertible.
Compute the unique solution using torch.linalg.solve(A,b).
Coefficient matrix A must be invertible.
X = torch.linalg.solve(A, b)
Display the solution.
Display the solution.
print("Solution:\n", X)
Check if the calculated solution is correct or not.
Check if the calculated solution is correct or not.
print(torch.allclose(A @ X, b))
# True for correct solution
Take a look at the following example −
# import required library
import torch

'''
Let's suppose our square system of linear equations is:
2x + 3y = 3
x - 2y = 0
'''

print("Linear equation:")
print("2x + 3y = 3")
print("x - 2y = 0")

# define the coefficient matrix A
A = torch.tensor([[2., 3.],[1., -2.]])
# define right hand side tensor b
b = torch.tensor([3., 0.])

# Solve the linear equation
X = torch.linalg.solve(A, b)

# print the solution of above linear equation
print("Solution:\n", X)

# check above solution to be true
print(torch.allclose(A @ X, b))
It will produce the following output −
Linear equation:
2x + 3y = 3
x - 2y = 0
Solution:
   tensor([0.8571, 0.4286])
True
Let's take another example −
# import required library
import torch

# define the coefficient matrix A for a 3x3
# square system of linear equations
A = torch.randn(3,3)

# define right hand side tensor b
b = torch.randn(3)

# Solve the linear equation
X = torch.linalg.solve(A, b)

# print the solution of above linear equation
print("Solution:\n", X)

# check above solution to be true
print(torch.allclose(A @ X, b))
It will produce the following output −
Solution:
   tensor([-0.2867, -0.9850, 0.9938])
True