Study Plan

Spectral Norm by Power Iteration

JAX

Medium

Estimate the largest singular value of a matrix with a fixed number of power iterations from a given start vector.

The matrix w has shape (M, N) and the start vector v0 has shape (N,), and $\lVert \cdot \rVert$ is the Euclidean length. Start from $v = v_0 / \lVert v_0 \rVert$ . Each of the num_iters iterations performs

u = \frac{w v}{\lVert w v \rVert}, \qquad v = \frac{w^\top u}{\lVert w^\top u \rVert}.

After the last iteration, using the final v, set $\sigma = \lVert w v \rVert$ and $u = w v / \sigma$ . With num_iters equal to 0, v is the normalized start vector.

Your function is compiled with jax.jit, and num_iters arrives as a traced int32 scalar, so Python's range cannot read it.

Return a dict with the key sigma, a float32 JAX array of shape (), the key u, a float32 JAX array of shape (M,), and the key v, a float32 JAX array of shape (N,).

Examples

Input: w = [[3.0,0.0],[0.0,1.0]], v0 = [1.0,1.0], num_iters = 2

Output: {"sigma":2.999797,"u":[0.999992,0.004115],"v":[0.999924,0.012345]}

Explanation: The unit start vector is [0.707107,0.707107]. Each iteration multiplies the ratio of the first component of v to the second by 9, so after two iterations v is close to [1,0] and sigma is close to the true spectral norm of 3.

Input: w = [[1.0,2.0,0.0],[2.0,1.0,1.0]], v0 = [1.0,0.0,0.0], num_iters = 0

Output: {"sigma":2.236068,"u":[0.447214,0.894427],"v":[1.0,0.0,0.0]}

Input: w = [[2.0,1.0],[1.0,2.0],[0.0,1.0]], v0 = [0.0,2.0], num_iters = 5

Output: {"sigma":3.087253,"u":[0.671542,0.700041,0.242847],"v":[0.661745,0.749729]}

Hint 1

jax.lax.fori_loop(0, num_iters, body, v0) with body(i, v) -> new_v.

Hint 2

jnp.linalg.norm(v) gives a vector's length as a scalar array.

Hint 3

W @ v and W.T @ u are the two halves of one iteration.

Requirements

Run the iterations with jax.lax.fori_loop so that the compiled program contains a while loop, which is what fori_loop produces when its upper bound is traced. A program without a while loop is rejected.
Keep the loop value at the same shape and dtype on every iteration.

Constraints

w is a float32 JAX array of shape (M, N) with M and N from 1 to 5.
v0 is a float32 JAX array of shape (N,) with a nonzero length.
num_iters is an int32 JAX scalar of shape () with a value from 0 to 20.
w v and the transpose of w times u never become the zero vector during the iteration.

Try Similar Problems

Traced Learning Rate ScheduleMedium Exponential Moving Average with ScanMedium Static Arguments in jitEasy Pad and Mask to Bucket LengthMedium Hessian-Vector ProductMedium Per-Example Gradient ClippingMedium Straight-Through Estimator with Custom VJPMedium Batched Attention with vmapMedium Pairwise Distances with Nested vmapMedium Global Norm Gradient ClippingMedium AdamW from ScratchMedium Mixed Precision with Loss ScalingMedium