Skip to content

← JAX Numerical Computing step 4 of 25

Easy Primitives

linalg.norm (Frobenius)

Jump to the task

Measure the combined size of matrix entries

The Frobenius norm is the square root of the sum of squared magnitudes of a matrix's entries. For a real matrix, the formula is sqrt(sum(A**2)). It is the same Euclidean length you would get by flattening all entries into one vector.

This gives one measure of the whole matrix, regardless of whether it is square. For example, the Frobenius norm of the difference between two matrices measures their overall entrywise discrepancy. Squared Frobenius norms also appear in weight penalties.

Worked example

import jax.numpy as jnp

A = jnp.array([[3.0, 0.0],
               [0.0, 4.0]])
frobenius = jnp.linalg.norm(A)             # 5.0
explicit = jnp.linalg.norm(A, ord="fro")   # 5.0
spectral = jnp.linalg.norm(A, ord=2)       # 4.0

The Frobenius calculation is sqrt(3**2 + 4**2) = 5. The spectral norm answers a different question: how much can the matrix stretch a unit vector? Here the largest stretch is four, along the second axis. More generally, the spectral norm is the largest singular value.

Problem

Implement matrix_frobenius(A) with jnp.linalg.norm.

  • A is a two-dimensional real float32 array of shape (m, n). It need not be square, and entries may be negative.
  • Return one scalar JAX floating-point value: the Frobenius norm.
  • An all-zero matrix has norm zero. A matrix with an empty dimension also has norm zero, since there are no entries to contribute.

For this two-dimensional input, either the default call or explicit ord="fro" computes the requested norm. Do not use ord=2, which would request the spectral norm instead.

Axes change what is measured

With no axis argument, this call combines all entries. Supplying axis=1 instead gives a vector norm for each row, not a single norm for the whole matrix.

For a batch of matrices with shape (B, m, n), use jnp.linalg.norm(batch, axis=(-2, -1)) to get one Frobenius norm per matrix. vmap is another option, but is not required. A default call without axes on that batch produces one norm across the entire array.

Values and gradients are different questions

The Frobenius norm is differentiable at nonzero real matrices, with gradient A / norm(A). At the all-zero matrix, the norm is defined but its derivative is not. Automatic differentiation can therefore produce NaNs there; a valid forward value does not imply a valid gradient. The squared norm has gradient 2*A, including at zero, when expressed directly as the sum of squared entries.

Like other float32 computations, the norm can lose precision or overflow for extreme values. Use tolerances when comparing computed results.