<p dir="auto">This appears to be valid:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import os
os.environ[&quot;XLA_FLAGS&quot;] = '--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)))"><pre class="notranslate"><code class="notranslate">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)))
</code></pre></div>
<p dir="auto">It appears that the non-diff arguments are brought to the front for <code class="notranslate">bwd</code>, but not <code class="notranslate">fwd</code>.<br>
However, the document implies that they should be moved for both <code class="notranslate">fwd</code> and <code class="notranslate">bwd</code>:<br>
<a href="https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html#jax.custom_vjp-with-nondiff_argnums" rel="nofollow">https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html#jax.custom_vjp-with-nondiff_argnums</a></p>
<p dir="auto">The <code class="notranslate">lax.custom_root</code> 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 <code class="notranslate">stop_gradient</code> <a href="https://github.com/google/jax/blob/e6950dd63ed3c054067c53bd4806dfbb85e0eef7/jax/_src/lax/control_flow.py#L2310">here</a>.</p>
<p dir="auto">While I've included a minimal example below, <a href="https://colab.research.google.com/drive/1e5_0gGYc-v9lKKwvUS7rqSUhA0D-OpT3" rel="nofollow">this notebook</a> does a deeper dive into the error and demonstrates that removing <code class="notranslate">_stop_gradient_fun</code> applications resolves the problem.</p>
<p dir="auto">I'd be happy to make a PR to resolve the issue, but my naïve solution of removing all calls to <code class="notranslate">stop_gradient</code> in <code class="notranslate">custom_root</code> appears... well, naïve. I'm just not sure why <code class="notranslate">stop_gradient</code> 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...</p>
<h3 dir="auto">Related issues</h3>
<ul dir="auto">
<li><a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="502922387" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/1448" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/1448/hovercard" href="https://github.com/google/jax/issues/1448">#1448</a> discusses a <code class="notranslate">define_implicit_gradient</code> function that doesn't use <code class="notranslate">stop_gradient</code>. This function was never merged. Regardless, the fact that <code class="notranslate">custom_root</code> methods can't use gradients seems like a bug.
<ul dir="auto">
<li>Additional reading of the discussion there suggests that one of the main reasons for <code class="notranslate">stop_gradient</code> was the non-differentiability of control structures (<code class="notranslate">lax.while_loop</code> in particular). Nevertheless, my tests don't seem to trigger the issue.</li>
</ul>
</li>
</ul>
<h2 dir="auto">Minimal example</h2>
<p dir="auto">This example follows <a href="https://github.com/google/jax/blob/e6950dd63ed3c054067c53bd4806dfbb85e0eef7/tests/lax_control_flow_test.py#L2010">the JAX unit test</a> of the <code class="notranslate">custom_root</code> method, with the bisection method replaced with Newton's method.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
import numpy as np

def newton(f, x0):
  &quot;&quot;&quot;Newton's method for root-finding.&quot;&quot;&quot;
  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 &lt; 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]"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">np</span>

<span class="pl-k">def</span> <span class="pl-en">newton</span>(<span class="pl-s1">f</span>, <span class="pl-s1">x0</span>):
  <span class="pl-s">"""Newton's method for root-finding."""</span>
  <span class="pl-s1">initial_state</span> <span class="pl-c1">=</span> (<span class="pl-c1">0</span>, <span class="pl-s1">x0</span>)  <span class="pl-c"># (iteration, x)</span>

  <span class="pl-k">def</span> <span class="pl-en">cond</span>(<span class="pl-s1">state</span>):
    <span class="pl-s1">it</span>, <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">state</span>
    <span class="pl-c"># We fix 10 iterations for simplicity, this is plenty for convergence in our tests.</span>
    <span class="pl-k">return</span> (<span class="pl-s1">it</span> <span class="pl-c1">&lt;</span> <span class="pl-c1">10</span>)

  <span class="pl-k">def</span> <span class="pl-en">body</span>(<span class="pl-s1">state</span>):
    <span class="pl-s1">it</span>, <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">state</span>
    <span class="pl-s1">fx</span>, <span class="pl-s1">dfx</span> <span class="pl-c1">=</span> <span class="pl-en">f</span>(<span class="pl-s1">x</span>), <span class="pl-s1">jax</span>.<span class="pl-en">grad</span>(<span class="pl-s1">f</span>)(<span class="pl-s1">x</span>)
    <span class="pl-s1">step</span> <span class="pl-c1">=</span> <span class="pl-s1">fx</span> <span class="pl-c1">/</span> <span class="pl-s1">dfx</span>
    <span class="pl-s1">new_state</span> <span class="pl-c1">=</span> <span class="pl-s1">it</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>, <span class="pl-s1">x</span> <span class="pl-c1">-</span> <span class="pl-s1">step</span>
    <span class="pl-k">return</span> <span class="pl-s1">new_state</span>

  <span class="pl-k">return</span> <span class="pl-s1">jax</span>.<span class="pl-s1">lax</span>.<span class="pl-en">while_loop</span>(
    <span class="pl-s1">cond</span>,
    <span class="pl-s1">body</span>,
    <span class="pl-s1">initial_state</span>,
  )[<span class="pl-c1">1</span>]</pre></div>
<h3 dir="auto">Newton's method without <code class="notranslate">custom_root</code></h3>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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(&quot;Without custom_root&quot;)
v = sqrt_cubed_newton(x_test)
print(f&quot;Value = {v}&quot;)
if abs(v - x_test ** 1.5) &lt; 1e-5:
  print(&quot;Pass&quot;)
else:
  print(&quot;Fail&quot;)"><pre class="notranslate"><span class="pl-s1">initial_guess</span> <span class="pl-c1">=</span> <span class="pl-c1">2.0</span>
<span class="pl-s1">x_test</span> <span class="pl-c1">=</span> <span class="pl-c1">5.0</span>

<span class="pl-k">def</span> <span class="pl-en">sqrt_cubed_newton</span>(<span class="pl-s1">x</span>):
  <span class="pl-s1">implicit</span> <span class="pl-c1">=</span> <span class="pl-k">lambda</span> <span class="pl-s1">y</span>: <span class="pl-s1">y</span> <span class="pl-c1">**</span> <span class="pl-c1">2</span> <span class="pl-c1">-</span> <span class="pl-s1">x</span> <span class="pl-c1">**</span> <span class="pl-c1">3</span>
  <span class="pl-k">return</span> <span class="pl-en">newton</span>(<span class="pl-s1">implicit</span>, <span class="pl-s1">initial_guess</span>)

<span class="pl-en">print</span>(<span class="pl-s">"Without custom_root"</span>)
<span class="pl-s1">v</span> <span class="pl-c1">=</span> <span class="pl-en">sqrt_cubed_newton</span>(<span class="pl-s1">x_test</span>)
<span class="pl-en">print</span>(<span class="pl-s">f"Value = <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">v</span><span class="pl-kos">}</span></span>"</span>)
<span class="pl-k">if</span> <span class="pl-en">abs</span>(<span class="pl-s1">v</span> <span class="pl-c1">-</span> <span class="pl-s1">x_test</span> <span class="pl-c1">**</span> <span class="pl-c1">1.5</span>) <span class="pl-c1">&lt;</span> <span class="pl-c1">1e-5</span>:
  <span class="pl-en">print</span>(<span class="pl-s">"Pass"</span>)
<span class="pl-k">else</span>:
  <span class="pl-en">print</span>(<span class="pl-s">"Fail"</span>)</pre></div>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Without custom_root
Value=11.180339813232422
Pass"><pre class="notranslate"><code class="notranslate">Without custom_root
Value=11.180339813232422
Pass
</code></pre></div>
<h3 dir="auto">With <code class="notranslate">custom_root</code></h3>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="
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(&quot;With custom_root&quot;)
v = sqrt_cubed_newton_custom_root(x_test)
print(f&quot;Value = {v}&quot;)
if abs(v - x_test ** 1.5) &lt; 1e-5:
  print(&quot;Pass&quot;)
else:
  print(&quot;Fail&quot;)"><pre class="notranslate"><span class="pl-k">def</span> <span class="pl-en">sqrt_cubed_newton_custom_root</span>(<span class="pl-s1">x</span>):
  <span class="pl-k">def</span> <span class="pl-en">scalar_solve</span>(<span class="pl-s1">g</span>, <span class="pl-s1">y</span>):
    <span class="pl-k">return</span> <span class="pl-s1">y</span> <span class="pl-c1">/</span> <span class="pl-en">g</span>(<span class="pl-c1">1.0</span>)
  <span class="pl-s1">implicit</span> <span class="pl-c1">=</span> <span class="pl-k">lambda</span> <span class="pl-s1">y</span>: <span class="pl-s1">y</span> <span class="pl-c1">**</span> <span class="pl-c1">2</span> <span class="pl-c1">-</span> <span class="pl-s1">x</span> <span class="pl-c1">**</span> <span class="pl-c1">3</span>
  <span class="pl-k">return</span> <span class="pl-s1">jax</span>.<span class="pl-s1">lax</span>.<span class="pl-en">custom_root</span>(<span class="pl-s1">implicit</span>, <span class="pl-s1">initial_guess</span>, <span class="pl-s1">newton</span>, <span class="pl-s1">scalar_solve</span>)

<span class="pl-en">print</span>(<span class="pl-s">"With custom_root"</span>)
<span class="pl-s1">v</span> <span class="pl-c1">=</span> <span class="pl-en">sqrt_cubed_newton_custom_root</span>(<span class="pl-s1">x_test</span>)
<span class="pl-en">print</span>(<span class="pl-s">f"Value = <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">v</span><span class="pl-kos">}</span></span>"</span>)
<span class="pl-k">if</span> <span class="pl-en">abs</span>(<span class="pl-s1">v</span> <span class="pl-c1">-</span> <span class="pl-s1">x_test</span> <span class="pl-c1">**</span> <span class="pl-c1">1.5</span>) <span class="pl-c1">&lt;</span> <span class="pl-c1">1e-5</span>:
  <span class="pl-en">print</span>(<span class="pl-s">"Pass"</span>)
<span class="pl-k">else</span>:
  <span class="pl-en">print</span>(<span class="pl-s">"Fail"</span>)</pre></div>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="With custom_root
Value = nan
Fail"><pre class="notranslate"><code class="notranslate">With custom_root
Value = nan
Fail
</code></pre></div>
<p dir="auto">The notebook referenced above shows how removing the call to <code class="notranslate">_stop_gradient_fun</code> fixes the problem.</p>
<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.</li>
</ul>