Skip to content

← JAX Numerical Computing step 1 of 25

Medium Primitives

linalg.solve

Jump to the task

Find the unknown vector

A linear system A @ x = b describes several equations sharing the same unknowns. The rows of A contain their coefficients, and b contains the desired results. jnp.linalg.solve(A, b) finds x.

When you need the solution vector, solve directly instead of first computing jnp.linalg.inv(A) @ b. Forming the full inverse does extra work and introduces another step where rounding can affect the result. Solving directly does not, however, remove sensitivity already present in the equations.

Worked example

import jax.numpy as jnp

A = jnp.array([[2.0, 1.0],
               [1.0, 3.0]])
b = jnp.array([4.0, 5.0])
x = jnp.linalg.solve(A, b)
# approximately [1.4, 1.2]
reconstructed = A @ x
# approximately [4.0, 5.0]

The first row means 2*x[0] + x[1] = 4; the second means x[0] + 3*x[1] = 5. Substituting 1.4 and 1.2 satisfies both equations. Checking A @ x against b is a useful way to check your result without knowing the solution in advance.

Problem

Implement solve_linear_system(A, b) using jnp.linalg.solve.

  • A is a square, nonsingular float32 array of shape (n, n).
  • b is a float32 vector of shape (n,).
  • Return a vector of shape (n,) satisfying A @ x approximately equal to b. Use floating-point tolerance when comparing results.

Nonsingular means the matrix has an inverse and the system has a unique solution for each right-hand side. No symmetry or positive-definiteness assumption is needed for this exercise.

What a solver cannot guarantee

A nearly singular matrix can amplify small changes in the inputs or rounding errors. This is called poor conditioning. A direct solve helps avoid unnecessary work, but it cannot guarantee an accurate solution for an ill-conditioned problem. Even a small residual A @ x - b does not by itself guarantee that x is close to the exact solution.

For singular matrices, JAX may return NaN or infinity rather than raise a LinAlgError. They are outside this exercise's contract. Do not add a diagonal term to make a matrix invertible unless you intend to change the equations being solved.

For rectangular systems, least squares is a different task handled by jnp.linalg.lstsq. For symmetric positive-definite matrices, a Cholesky factorization can be useful; jax.scipy.linalg.cho_solve consumes an existing Cholesky factor and its triangle flag. Positive semidefinite matrices can be singular, so that weaker condition is not enough.