Skip to content

← JAX Numerical Computing step 7 of 25

Medium Primitives

QR Decomposition

Jump to the task

Separate directions from coefficients

QR decomposition writes a real matrix as A = Q @ R. The columns of Q are perpendicular unit vectors. R holds the coefficients needed to combine those vectors into the columns of A.

In JAX's default reduced mode, let k = min(m, n) for an input with shape (m, n). Then Q has shape (m, k) and R has shape (k, n). Entries of R below its main diagonal are zero. For a wide input, R is rectangular; for a tall or square input, it is square.

Worked example

import jax.numpy as jnp

A = jnp.array([[1.0, 0.0],
               [0.0, 1.0],
               [1.0, 1.0]])
Q, R = jnp.linalg.qr(A)
# Q.shape is (3, 2); R.shape is (2, 2)

reconstructed = Q @ R              # approximately A
inner_products = Q.T @ Q           # approximately the 2x2 identity
q_norm = jnp.linalg.norm(Q)         # approximately 1.4142

Each of Q's two columns has squared length one. The sum of the squares of all its entries is therefore two, so its Frobenius norm is sqrt(2). This is the matrix norm that jnp.linalg.norm(Q) computes by default.

Problem

Implement qr_decompose_q_norm(A) by computing a reduced QR decomposition and returning the Frobenius norm of Q.

  • A is a two-dimensional real float32 array of shape (m, n).
  • Use jnp.linalg.qr(A), whose default mode is "reduced".
  • Return a scalar JAX value, not Q, R, or a pair of arrays.

The expected norm is approximately sqrt(min(m, n)), because each column contributes one to its square. This identity is a useful check on your result; practice computing the factors in your implementation. Even a rank-deficient or zero matrix gets orthonormal columns in Q. If a dimension is empty, Q has no columns and its norm is zero.

Interpreting the factors

With independent columns and m >= n, reduced Q is an orthonormal basis for A's column space. If A is rank deficient, Q includes extra directions beyond that space. Its column count, and thus its norm, does not tell you A's rank.

The factors are not unique: negating a column of Q and the corresponding row of R preserves their product. Check reconstruction and orthogonality rather than expecting particular signs. Rank-deficient inputs can have additional freedom in the choice of directions.

mode="complete" returns a square Q of shape (m, m). For a tall input this adds columns and changes Q's norm to sqrt(m), so it is not the mode requested here.

How QR helps with least squares

For a tall matrix with independent columns, minimizing ||A @ x - b|| reduces to solving the triangular system R @ x = Q.T @ b. This avoids forming the normal-equation matrix A.T @ A, which squares the 2-norm condition number. It does not remove sensitivity already present in A. If columns are dependent, that triangular system is singular and needs a method that handles rank deficiency, such as an SVD-based solver.