Skip to content

← JAX Numerical Computing step 8 of 25

Medium Primitives

SVD Singular Values

Jump to the task

Measure a matrix's stretches

Singular values describe the scale factors in a singular value decomposition, or SVD. They are nonnegative and returned from largest to smallest. Unlike an eigendecomposition, SVD works for rectangular matrices as well as square ones.

The largest singular value is the matrix's spectral norm: its largest stretch of a unit vector. Smaller singular values describe the other directions in the decomposition. For shape (m, n), JAX returns min(m, n) singular values.

Worked example

import jax.numpy as jnp

A = jnp.array([[-3.0, 0.0, 0.0],
               [0.0, 4.0, 0.0]])
s = jnp.linalg.svd(A, compute_uv=False)
# approximately [4.0, 3.0]

This matrix scales one input direction by magnitude three and another by four, while discarding the third coordinate. Its negative entry changes a direction, not the sign of a singular value. The result has two entries because min(2, 3) = 2; the wide matrix's extra null-space direction does not add another entry to this returned vector.

Where the values fit in the decomposition

If you request vectors too, JAX returns (U, s, Vh). For a real matrix, the reduced decomposition uses full_matrices=False and has shapes:

  • U: (m, k)
  • s: (k,)
  • Vh: (k, n), where k = min(m, n)

Reconstruct it as U @ jnp.diag(s) @ Vh. JAX returns the transposed right-vector matrix already; do not transpose Vh again. For complex matrices it is the conjugate transpose, though this exercise uses real inputs. Default full matrices have different shapes, so the reduced reconstruction formula should not be applied blindly to them.

Problem

Implement svd_singular_values(A) with jnp.linalg.svd(A, compute_uv=False).

  • A is a two-dimensional real float32 array of shape (m, n).
  • Return a one-dimensional array of shape (min(m, n),).
  • Keep JAX's descending order, including any zero or near-zero values. Do not return the singular vectors or discard small values.

An all-zero matrix has all-zero singular values. If a dimension is empty, the returned singular-value array is empty.

Reading the result

In exact arithmetic, the number of nonzero singular values equals the rank. Numerically, a mathematically zero value may appear as a tiny positive number. Rank estimation therefore needs a tolerance chosen with the scale and precision in mind, not simply a count of s != 0. This exercise returns the computed values without estimating rank.

For a nonsingular square matrix, s[0] / s[-1] is its condition number in the 2-norm. A large ratio indicates sensitivity; a zero smallest singular value means singularity. Do not interpret tiny roundoff values as proof of invertibility.

Low-rank approximation keeps the leading singular values and their corresponding vectors. Singular values alone cannot reconstruct the matrix. For this task, compute_uv=False requests only what you need.