Skip to content

← JAX Numerical Computing step 9 of 25

Hard Primitives

eigh: Symmetric Eigendecomp

Jump to the task

Find a matrix's eigenvalues

An eigenvector is a nonzero vector that a matrix maps to a scalar multiple of itself: A @ v = eigenvalue * v. The eigenvalue is that scale factor. A negative factor reverses the direction; a zero factor maps it to zero.

A real symmetric matrix has real eigenvalues and can be decomposed using an orthonormal basis of eigenvectors. jnp.linalg.eigh(A) returns the eigenvalues in ascending order and the corresponding eigenvectors as columns. Symmetry lets the routine use methods designed for this structure.

For example, a covariance matrix in PCA is symmetric. Its eigenvectors identify principal directions, and its eigenvalues describe variance along them. Other symmetric matrices can have negative eigenvalues; positive definiteness is not required by this exercise.

Worked example

import jax.numpy as jnp

A = jnp.array([[3.0, 1.0],
               [1.0, 3.0]])
eigenvalues, eigenvectors = jnp.linalg.eigh(A)
# eigenvalues: approximately [2.0, 4.0]
largest = eigenvalues[-1]
# approximately 4.0

The vector [1, 1] becomes [4, 4], so it has eigenvalue 4. The vector [1, -1] becomes [2, -2], so it has eigenvalue 2. The returned eigenvectors are normalized versions of these directions. Their signs can differ without changing the decomposition.

Problem

Implement eigh_largest(A) using jnp.linalg.eigh and return the last eigenvalue.

  • A is a real symmetric float32 array of shape (n, n), with n >= 1.
  • Return one scalar JAX floating-point value.
  • Largest means greatest numerical value, not greatest absolute value. For eigenvalues [-7, -2], return -2, not 7 or -7.

Do not take absolute values, clamp negative eigenvalues, or return an eigenvector. Singular symmetric matrices are allowed; zero can be an eigenvalue. Compare computed values with a floating-point tolerance.

Input symmetry and numerical limits

JAX symmetrizes the input by default using both triangles. For real matrices this means (A + A.T) / 2. An asymmetric input would therefore describe a different matrix from the one you supplied. This exercise provides symmetric inputs; no repair is requested.

For a general nonsymmetric matrix, jnp.linalg.eig is a different API and eigenvalues may be complex. Do not use eigh merely to force such results to be real. Complex Hermitian matrices are supported by the general eigh API but are outside this real-valued exercise.

If an eigenvalue is repeated, there is no unique eigenvector basis for its eigenspace. Check eigenvector equations or reconstructed matrices rather than expecting identical vector entries from every backend. Floating-point accuracy and elapsed time depend on the input and system; there is no fixed two-times accuracy or speed improvement to promise.