Issue 1:
This appears to be valid:

import os
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'

import functools

import jax
import jax.numpy as jnp
import numpy as np

@functools.partial(jax.custom_vjp, nondiff_argnums=(1,))
def custom_psum(x, axis_name):
return x

def fwd(x, axis_name):
print('fwd_arg0', x)
print('fwd_arg1', axis_name)
return custom_psum(x, axis_name), None

def bwd(axis_name, res, g):
print('bwd_arg0', axis_name)
print('bwd_arg1', res)
print('bwd_arg2', g)
del axis_name, res
return (g,)

custom_psum.defvjp(fwd, bwd)

def loss_fn(x):
return jnp.sum(custom_psum(x, 'i'))

jax.pmap(jax.grad(loss_fn), 'i')(np.random.random((8, 2, 3)))

It appears that the non-diff arguments are brought to the front for bwd, but not fwd.

However, the document implies that they should be moved for both fwd and bwd:

https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html#jax.custom_vjp-with-nondiff_argnums

Issue 2:
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.

While I've included a minimal example below, this notebook 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 of the custom_root method, with the bisection method replaced with Newton's method.

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

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

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:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this

If applicable, include full error messages/tracebacks.