Repository: google/jax

Issue:
`lax.custom_root` fails on root finders that require gradient
The `lax.custom_root` method does not work on methods that require gradient, e.g., the Newton-Raphson method. The core of the problem seems to be overly-aggressive application of `stop_gradient` [here](https://github.com/google/jax/blob/e6950dd63ed3c054067c53bd4806dfbb85e0eef7/jax/_src/lax/control_flow.py#L2310).

While I've included a minimal example below, [this notebook](https://colab.research.google.com/drive/1e5_0gGYc-v9lKKwvUS7rqSUhA0D-OpT3) does a deeper dive into the error and demonstrates that removing `_stop_gradient_fun` applications resolves the problem.

I'd be happy to make a PR to resolve the issue, but my naïve solution of removing all calls to `stop_gradient` in `custom_root` appears... well, naïve. I'm just not sure why `stop_gradient` is required at all here, but it seems related to not wanting to trace gradients through control structures like while-loops. I suppose I'm curious why custom JVP'ing a function doesn't automatiically shortcut the gradient computation...

### Related issues

 - #1448 discusses a `define_implicit_gradient` function that doesn't use `stop_gradient`. This function was never merged. Regardless, the fact that `custom_root` methods can't use gradients seems like a bug.
   - Additional reading of the discussion there suggests that one of the main reasons for `stop_gradient` was the non-differentiability of control structures (`lax.while_loop` in particular). Nevertheless, my tests don't seem to trigger the issue.


## Minimal example

This example follows [the JAX unit test](https://github.com/google/jax/blob/e6950dd63ed3c054067c53bd4806dfbb85e0eef7/tests/lax_control_flow_test.py#L2010) of the `custom_root` method, with the bisection method replaced with Newton's method.


```python
import jax
import numpy as np

def newton(f, x0):
  """Newton's method for root-finding."""
  initial_state = (0, x0)  # (iteration, x)

  def cond(state):
    it, x = state
    # We fix 10 iterations for simplicity, this is plenty for convergence in our tests.
    return (it < 10)

  def body(state):
    it, x = state
    fx, dfx = f(x), jax.grad(f)(x)
    step = fx / dfx
    new_state = it + 1, x - step
    return new_state

  return jax.lax.while_loop(
    cond,
    body,
    initial_state,
  )[1]
```


### Newton's method without `custom_root`

```python
initial_guess = 2.0
x_test = 5.0

def sqrt_cubed_newton(x):
  implicit = lambda y: y ** 2 - x ** 3
  return newton(implicit, initial_guess)

print("Without custom_root")
v = sqrt_cubed_newton(x_test)
print(f"Value = {v}")
if abs(v - x_test ** 1.5) < 1e-5:
  print("Pass")
else:
  print("Fail")
```

```
Without custom_root
Value=11.180339813232422
Pass
```

### With `custom_root`

```python

def sqrt_cubed_newton_custom_root(x):
  def scalar_solve(g, y):
    return y / g(1.0)
  implicit = lambda y: y ** 2 - x ** 3
  return jax.lax.custom_root(implicit, initial_guess, newton, scalar_solve)

print("With custom_root")
v = sqrt_cubed_newton_custom_root(x_test)
print(f"Value = {v}")
if abs(v - x_test ** 1.5) < 1e-5:
  print("Pass")
else:
  print("Fail")
```
```
With custom_root
Value = nan
Fail
```

The notebook referenced above shows how removing the call to `_stop_gradient_fun` fixes the problem.


Please:

- [x] Check for duplicate issues.
- [x] Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this
- [x] If applicable, include full error messages/tracebacks.



Discussion:
Hi @mbmccoy, thanks for the careful investigation here!

I’m struggling to remember the exact context here, but I think we might actually be able to get away with removing the use of stop_gradient inside custom_root, because while_loop now has a JVP rule defined. This was not the case back when custom_root was written. 
We also switched from writing custom_root as a primitive to writing it via custom_jvp. So that also may have obviated the need for this.

Basically, I would suggest going ahead with removing stop_gradient_fun, being sure to add something like your Newton solver as an integration test.

Resolving patch:
diff --git a/jax/_src/lax/control_flow.py b/jax/_src/lax/control_flow.py
--- a/jax/_src/lax/control_flow.py
+++ b/jax/_src/lax/control_flow.py
@@ -2357,18 +2357,6 @@ def _promote_weak_typed_inputs(in_vals, in_avals, out_avals):
     in_vals[i] = lax.convert_element_type(in_vals[i], new_dtype)
   return in_vals, True
 
-def _stop_gradient_fun(f):
-  """Create a version of f() that stops all gradients."""
-  def wrapper(*args, **kwargs):
-    args_flat, in_args_tree = tree_flatten((args, kwargs))
-    args_avals = tuple(_map(_abstractify, args_flat))
-    g = lambda a, b: f(*a, **b)
-    jaxpr, consts, out_tree = _initial_style_jaxpr(g, in_args_tree, args_avals)
-    all_args = _map(lax.stop_gradient, (*consts, *args_flat))
-    out = core.jaxpr_as_fun(jaxpr)(*all_args)
-    return tree_unflatten(out_tree, out)
-  return wrapper
-
 
 _RootTuple = collections.namedtuple('_RootTuple', 'f, solve, l_and_s')
 
@@ -2426,7 +2414,7 @@ def custom_root(f, initial_guess, solve, tangent_solve, has_aux=False):
   _check_tree("f", "initial_guess", out_tree, in_tree, False)
 
   solve_jaxpr, solve_consts, solution_tree = _initial_style_jaxpr(
-      partial(solve, _stop_gradient_fun(f)), in_args_tree, guess_avals)
+      partial(solve, f), in_args_tree, guess_avals)
   _check_tree("solve", "initial_guess", solution_tree, in_tree, has_aux)
 
   def linearize_and_solve(x, b):


Test patch:
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -95,6 +95,52 @@ def posify(matrix):
   jtu.ignore_warning, message=".*jit-of-pmap.*")
 
 
+# Simple optimization routine for testing custom_root
+def binary_search(func, x0, low=0.0, high=100.0):
+  del x0  # unused
+
+  def cond(state):
+    low, high = state
+    midpoint = 0.5 * (low + high)
+    return (low < midpoint) & (midpoint < high)
+
+  def body(state):
+    low, high = state
+    midpoint = 0.5 * (low + high)
+    update_upper = func(midpoint) > 0
+    low = jnp.where(update_upper, low, midpoint)
+    high = jnp.where(update_upper, midpoint, high)
+    return (low, high)
+
+  solution, _ = lax.while_loop(cond, body, (low, high))
+  return solution
+
+# Optimization routine for testing custom_root.
+def newton_raphson(func, x0):
+  tol = 1e-16
+  max_it = 20
+
+  fx0, dfx0 = func(x0), jax.jacobian(func)(x0)
+  initial_state = (0, x0, fx0, dfx0)  # (iteration, x, f(x), grad(f)(x))
+
+  def cond(state):
+    it, _, fx, _ = state
+    return (jnp.max(jnp.abs(fx)) > tol) & (it < max_it)
+
+  def body(state):
+    it, x, fx, dfx = state
+    step = jnp.linalg.solve(
+      dfx.reshape((-1, fx.size)), fx.ravel()
+    ).reshape(fx.shape)
+    x_next = x - step
+    fx, dfx = func(x_next), jax.jacobian(func)(x_next)
+    return (it + 1, x_next, fx, dfx)
+
+  _, x, _, _ = lax.while_loop(cond, body, initial_state)
+
+  return x
+
+
 class LaxControlFlowTest(jtu.JaxTestCase):
 
   def setUp(self):
@@ -2007,33 +2053,19 @@ def fun(carry, _):
 
     jax.grad(lambda x: jit_run_scan(x))(0.)  # doesn't crash
 
-  def test_custom_root_scalar(self):
+  @parameterized.named_parameters(
+      {"testcase_name": "binary_search", "solve_method": binary_search},
+      {"testcase_name": "newton_raphson", "solve_method": newton_raphson},
+  )
+  def test_custom_root_scalar(self, solve_method):
 
     def scalar_solve(f, y):
       return y / f(1.0)
 
-    def binary_search(func, x0, low=0.0, high=100.0):
-      del x0  # unused
-
-      def cond(state):
-        low, high = state
-        midpoint = 0.5 * (low + high)
-        return (low < midpoint) & (midpoint < high)
-
-      def body(state):
-        low, high = state
-        midpoint = 0.5 * (low + high)
-        update_upper = func(midpoint) > 0
-        low = jnp.where(update_upper, low, midpoint)
-        high = jnp.where(update_upper, midpoint, high)
-        return (low, high)
-
-      solution, _ = lax.while_loop(cond, body, (low, high))
-      return solution
-
     def sqrt_cubed(x, tangent_solve=scalar_solve):
       f = lambda y: y ** 2 - x ** 3
-      return lax.custom_root(f, 0.0, binary_search, tangent_solve)
+      # Note: Nonzero derivative at x0 required for newton_raphson
+      return lax.custom_root(f, 1.0, solve_method, tangent_solve)
 
     value, grad = jax.value_and_grad(sqrt_cubed)(5.0)
     self.assertAllClose(value, 5 ** 1.5, check_dtypes=False, rtol=1e-6)
@@ -2044,7 +2076,11 @@ def sqrt_cubed(x, tangent_solve=scalar_solve):
 
     inputs = jnp.array([4.0, 5.0])
     results = jax.vmap(sqrt_cubed)(inputs)
-    self.assertAllClose(results, inputs ** 1.5, check_dtypes=False)
+    self.assertAllClose(
+      results, inputs ** 1.5, check_dtypes=False,
+      atol={jnp.float32: 1e-3, jnp.float64: 1e-6},
+      rtol={jnp.float32: 1e-3, jnp.float64: 1e-6},
+    )
 
     results = jax.jit(sqrt_cubed)(5.0)
     self.assertAllClose(
@@ -2073,6 +2109,30 @@ def linear_solve(a, b):
     expected = jnp.linalg.solve(a, b)
     self.assertAllClose(expected, actual)
 
+  def test_custom_root_vector_nonlinear(self):
+
+    def nonlinear_func(x, y):
+      # func(x, y) == 0 if and only if x == y.
+      return (x - y) * (x**2 + y**2 + 1)
+
+    def tangent_solve(g, y):
+      return jnp.linalg.solve(
+        jax.jacobian(g)(y).reshape(-1, y.size),
+        y.ravel()
+      ).reshape(y.shape)
+
+    def nonlinear_solve(y):
+      f = lambda x: nonlinear_func(x, y)
+      x0 = -jnp.ones_like(y)
+      return lax.custom_root(f, x0, newton_raphson, tangent_solve)
+
+    y = self.rng().randn(3, 1)
+    jtu.check_grads(nonlinear_solve, (y,), order=2,
+                rtol={jnp.float32: 1e-2, jnp.float64: 1e-3})
+
+    actual = jax.jit(nonlinear_solve)(y)
+    self.assertAllClose(y, actual, rtol=1e-5, atol=1e-5)
+
   def test_custom_root_with_custom_linear_solve(self):
 
     def linear_solve(a, b):