Skip to content

← JAX Autodiff step 12 of 25

Hard Primitives

Custom VJP: Implicit Function Theorem

Jump to the task

A solver can define a quantity through an equation rather than an explicit formula. Implicit differentiation derives how the solution changes from that equation, without differentiating every solver step. It requires a sufficiently accurate solution and suitable local regularity conditions; it is not a guarantee that an arbitrary solver output has the right derivative.

Derive the square-root rule

The positive square root satisfies y ** 2 - x = 0. Differentiate this equation with respect to x:

2 * y * dy/dx - 1 = 0
dy/dx = 1 / (2 * y)

For x > 0, the positive root has y > 0, so division by 2 * y is valid. At zero, this finite derivative formula no longer applies. For coupled vector equations, implicit differentiation generally requires solving a linear system, not just elementwise division.

Approximate the root with Newton steps

Starting with y = 1, repeat:

y = 0.5 * (y + x / y)

This exercise uses exactly ten iterations. For the moderate positive inputs in the examples, they approximate the positive square root closely. For [4, 9], the outputs are approximately [2, 3].

Ten steps are not a universal convergence guarantee. For a very small input, such as 1e-12, the iterate remains near 0.0009766 after ten steps, far above the true root 0.000001. Very large inputs can also need more iterations. A production solver needs an appropriate initial guess, numerical safeguards and a convergence test.

Save the iterate for backward

Decorate the forward solver with jax.custom_vjp. Its forward rule returns (y, y): the computed output and the same value as a residual. The backward rule returns:

return (cotangent * 0.5 / y,)

Both y and the incoming cotangent are vectors, so multiplication and division are elementwise. The one-element tuple corresponds to the primal function's single argument.

This rule evaluates the implicit derivative formula at the saved approximation. It is not, in general, the exact derivative of ten Newton iterations. When the solve has converged closely, it agrees closely with the root derivative. When it has not, custom VJP does not repair that discrepancy. Differentiating the finite iteration itself is also valid, but answers a different question.

Saving the final iterate avoids replaying this solver in the specified first-order backward rule. That alone does not prove constant total memory or faster execution for an arbitrary implicit model.

Problem

Implement fixed_point_sqrt(x) with this specified approximation:

  1. Start from jnp.ones_like(x) and perform exactly ten Newton updates.
  2. Return the final iterate as the primal output.
  3. In the custom forward rule, return that output and save it as residual.
  4. Register a backward rule returning (cotangent * 0.5 / y,).

x is a finite positive float32 JAX vector of shape (n,), with n >= 1. Leave it unchanged. Return a JAX vector of the same shape. Do not replace the required fixed-step solver with jnp.sqrt or differentiate the loop to define its custom backward rule.

The exercise's examples use inputs where ten steps are accurate to numerical tolerance. The procedure does not promise an accurate square root for every positive float32 value. Output tests also cannot prove that the requested custom derivative was registered; check a VJP with a nonuniform cotangent locally.