Skip to content

← JAX Vectorization & Control Flow step 10 of 25

Medium Primitives

Newton's Method via lax.while_loop

Jump to the task

Refine a square-root estimate

To estimate sqrt(x), look for a positive y satisfying y*y - x = 0. Newton's method uses the derivative 2*y to update the estimate:

new_y = y - (y*y - x)/(2*y)
      = 0.5 * (y + x/y)

Start at y=1. For x=4, the estimates are approximately 2.5, 2.05, 2.000610 and 2.000000. Their residuals abs(y*y - 4) are approximately 2.25, 0.2025, 0.002439 and less than 0.000001. A tolerance of 0.001 therefore requires four updates.

The tolerance measures the residual abs(y*y-x), not the error abs(y - sqrt(x)) or the change between successive estimates. These are different stopping criteria. This exercise uses the first one.

Carry the count and estimate

Use (iteration_count, y) as the state of lax.while_loop. The condition accepts that state and returns one scalar boolean. The body accepts the state and returns a new pair with the same structure, shapes and dtypes.

Continue while the residual is at least tol and fewer than 100 updates have run. Each body call increments the count and applies one Newton update. If the initial y already has residual below tol, no updates are needed.

A runtime residual cannot control a Python while statement inside JIT: Python would need a concrete boolean during tracing. A lax.while_loop evaluates its condition inside the compiled loop. Its runtime iteration count does not support ordinary reverse-mode differentiation through the loop.

Limit the work when rounding prevents convergence

In float32, the estimate for sqrt(2) can settle at about 1.4142135 while its residual is still nonzero. With tol=1e-12, further updates cannot meet the requested tolerance. A residual-only loop can then run forever, despite the estimate already being useful.

The 100-update cap guarantees a bounded number of updates. It does not guarantee convergence. Return the latest estimate at the cap; callers who need a convergence decision must also check the residual. For production numerical work, use an appropriate library routine or a solver that reports its stopping status.

Problem

Implement while_sqrt(x, tol) using lax.while_loop, with initial state (0, 1.0). Inputs are positive finite float32 scalars. Continue while both iteration_count < 100 and jnp.abs(y*y - x) >= tol. Update to (iteration_count + 1, 0.5*(y + x/y)).

Return the final scalar y, whether the loop stopped because it met the tolerance or reached the cap. Do not return the count or replace the iteration with jnp.sqrt. For x=4 and tol=4, return the initial estimate 1: its residual 3 is already below 4. For x=4 and tol=0.001, return approximately 2 after four updates.

Use & with parenthesized scalar comparisons in the condition; Python and cannot decide a traced boolean. Output tests check the returned values, but cannot enforce the chosen loop primitive.