Skip to content

← JAX Numerical Computing step 6 of 25

Medium Primitives

Cholesky Decomposition

Jump to the task

Factor a matrix into triangular pieces

For a real symmetric positive-definite matrix A, Cholesky factorization finds a lower-triangular matrix L such that A = L @ L.T in exact arithmetic. Lower triangular means every entry above the diagonal is zero. The diagonal entries of L are positive.

Positive definite means v.T @ A @ v > 0 for every nonzero vector v. Positive semidefinite only requires this quantity to be nonnegative. That weaker condition permits singular matrices and is not sufficient for the standard Cholesky routine used here.

A Cholesky factor is useful when solving systems with the same matrix: solve L @ z = b, then L.T @ x = z. It also transforms independent standard normal samples into correlated samples: if z has covariance I, then L @ z has covariance L @ L.T, which is A.

Worked example

import jax.numpy as jnp

A = jnp.array([[4.0, 2.0],
               [2.0, 5.0]])
L = jnp.linalg.cholesky(A)
# [[2.0, 0.0],
#  [1.0, 2.0]]
reconstructed = L @ L.T
# approximately [[4.0, 2.0], [2.0, 5.0]]

To check the multiplication, the first diagonal entry is 2**2 = 4. The off-diagonal entries are 2*1 = 2, and the second diagonal entry is 1**2 + 2**2 = 5. This is a matrix factorization, not an elementwise square root of A.

Problem

Implement cholesky_factor(A) using jnp.linalg.cholesky with its default lower-triangular output.

  • A is a real float32 array of shape (n, n), symmetric and positive definite.
  • Return L of shape (n, n), with zeros above its diagonal and positive entries on the diagonal.
  • L @ L.T should agree with A within floating-point tolerance.

Things to watch

A singular positive-semidefinite matrix such as diag([1, 0]) does not meet the requirement. When factorization fails, JAX returns NaNs rather than a useful factor. Poor conditioning can also make the calculation sensitive to rounding, even for a mathematically positive-definite matrix.

By default, JAX symmetrizes its input using both triangles before factorization. Do not assume that only the lower triangle matters. This exercise supplies a symmetric matrix, so no repair or diagonal adjustment is requested. Adding a diagonal term changes the matrix whose factor you compute.

Cholesky takes advantage of the matrix's structure, but its elapsed time depends on matrix size, hardware, and implementation. There is no universal two-times speedup to rely on.