Skip to content

← JAX Numerical Computing step 2 of 25

Easy Primitives

linalg.inv

Jump to the task

Compute the matrix that undoes a transformation

For a square nonsingular matrix A, its inverse A_inv satisfies A @ A_inv = I and A_inv @ A = I in exact arithmetic. I is the identity matrix: multiplying by it leaves a vector unchanged.

jnp.linalg.inv(A) returns the entire inverse. This exercise asks for that matrix. If you only need a vector satisfying A @ x = b, use jnp.linalg.solve(A, b) instead of forming the inverse and multiplying. An inverse appearing in a mathematical formula does not automatically mean it should be explicitly computed in code.

Worked example

import jax.numpy as jnp

A = jnp.array([[2.0, 0.0],
               [0.0, 4.0]])
A_inv = jnp.linalg.inv(A)
# [[0.5, 0.0],
#  [0.0, 0.25]]
product = A @ A_inv
# approximately [[1.0, 0.0], [0.0, 1.0]]

This diagonal matrix doubles the first coordinate and quadruples the second. Its inverse halves the first coordinate and quarters the second. Taking elementwise reciprocals is not a general method for inverting a matrix; this shortcut applies only to its nonzero diagonal entries when the off-diagonal entries are zero.

Problem

Implement matrix_inverse(A) using jnp.linalg.inv.

  • A is a square, nonsingular float32 array of shape (n, n).
  • Return the inverse with shape (n, n).
  • Products with the original matrix should approximate the identity. Compare with a tolerance rather than expecting exact equality.

For example, the inverse of [[1, 2], [3, 4]] is approximately [[-2, 1], [1.5, -0.5]].

Numerical limits

A singular matrix has no inverse. JAX may return NaNs or infinities for it rather than raise a LinAlgError; do not rely on an exception as your validity check. This exercise supplies nonsingular matrices.

An ill-conditioned matrix is sensitive to small changes in its entries, so its computed inverse may have substantial error even when all values are finite. The condition number helps assess that sensitivity; computing it does not repair the matrix or guarantee accuracy.

Adding a small diagonal term changes the matrix and therefore changes the requested inverse. Regularization can be appropriate in a model that calls for it, but it is not a silent fix for this exercise. Likewise, a pseudoinverse solves a different problem and is not requested here.

A dense inverse typically takes work proportional to n**3 and stores n**2 entries. Avoid computing it when a direct solve provides the quantity you actually need.