Repository: google/jax

Issue:
Forward-mode differentiation rule for 'custom_lin' not implemented
Running this code produces the above error.

```
@partial(jax.custom_vjp, nondiff_argnums=(0, 1))
def _clip_gradient(lo, hi, x):
  return x  # identity function

def clip_gradient_fwd(lo, hi, x):
    # return x, None
    return x, (hi, ) 

def clip_gradient_bwd(lo, hi, _, g):
    return (np.clip(g, lo, hi),)

_clip_gradient.defvjp(clip_gradient_fwd, clip_gradient_bwd)

def clip_gradient(x):
    lo = -1
    hi = x + 1  # causes things to break
    return _clip_gradient(lo, hi, x)

print(jax.grad(clip_gradient)(1.))
```

Replacing the residual with `None` (see commented out line in `clip_gradient_fwd`) makes the output

```
Traced<ConcreteArray(1.0)>with<JVPTrace(level=1/0)>
  with primal = DeviceArray(1., dtype=float32)
       tangent = Traced<ShapedArray(float32[]):JaxprTrace(level=0/0)>
```

I was also able to get a mismatched tracer levels  error on a more complex example.

To my understanding setting `hi = x + 1` is the issue as it creates a trace of `hi` with `x`.

This example may seem a bit contrived, but I originally came across this trying to set an initial step size for `odeint` (see #2604).

IIUC we want to make concrete all `JVPTracer` instances that are in static args [here](https://github.com/google/jax/blob/master/jax/custom_derivatives.py#L441).


Discussion:
I also just managed to produce ```NotImplementedError: Batching rule for 'custom_lin' not implemented``` while doing something else.
As a workaround, you can write

```python
def clip_gradient(x):
    lo = -1
    hi = jax.lax.stop_gradient(x + 1)
    return _clip_gradient(lo, hi, x)
```

Resolving patch:
diff --git a/jax/custom_derivatives.py b/jax/custom_derivatives.py
--- a/jax/custom_derivatives.py
+++ b/jax/custom_derivatives.py
@@ -18,6 +18,7 @@
 import itertools as it
 import operator as op
 
+import jax
 from . import core
 from . import linear_util as lu
 from .tree_util import tree_flatten, tree_unflatten, tree_map, tree_multimap
@@ -82,6 +83,15 @@ def sum_tangents(x, *xs):
 def zeros_like_pytree(x):
   return tree_map(lambda _: zero, x)
 
+def stop_gradient(x):
+  return tree_map(_stop_gradient, x)
+
+def _stop_gradient(x):
+  if isinstance(x, core.Tracer) or core.valid_jaxtype(x):
+    return jax.lax.stop_gradient(x)
+  else:
+    return x
+
 
 ### JVPs
 
@@ -199,7 +209,10 @@ def __call__(self, *args, **kwargs):
       raise AttributeError(msg.format(self.__name__))
     args = _resolve_kwargs(self.fun, args, kwargs)
     if self.nondiff_argnums:
-      dyn_argnums = [i for i in range(len(args)) if i not in self.nondiff_argnums]
+      is_nondiff = [False] * len(args)
+      for i in self.nondiff_argnums: is_nondiff[i] = True
+      args = [stop_gradient(x) if b else x for b, x in zip(is_nondiff, args)]
+      dyn_argnums = [i for i, b in enumerate(is_nondiff) if not b]
       f_, dyn_args = argnums_partial(lu.wrap_init(self.fun), dyn_argnums, args)
       static_args = [args[i] for i in self.nondiff_argnums]
       jvp = _add_args(lu.wrap_init(self.jvp), static_args, left=True)
@@ -436,7 +449,10 @@ def __call__(self, *args, **kwargs):
       raise AttributeError(msg.format(self.__name__))
     args = _resolve_kwargs(self.fun, args, kwargs)
     if self.nondiff_argnums:
-      dyn_argnums = [i for i in range(len(args)) if i not in self.nondiff_argnums]
+      is_nondiff = [False] * len(args)
+      for i in self.nondiff_argnums: is_nondiff[i] = True
+      args = [stop_gradient(x) if b else x for b, x in zip(is_nondiff, args)]
+      dyn_argnums = [i for i, b in enumerate(is_nondiff) if not b]
       f_, dyn_args = argnums_partial(lu.wrap_init(self.fun), dyn_argnums, args)
       static_args = [args[i] for i in self.nondiff_argnums]
       fwd, _ = argnums_partial(lu.wrap_init(self.fwd), dyn_argnums, args)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2853,6 +2853,28 @@ def g(f, x):
 
     jax.grad(g, argnums=(1,))(F(2.0), 0.)  # doesn't crash
 
+  def test_nondiff_argnums_stop_gradient(self):
+    # https://github.com/google/jax/issues/2784
+    @partial(api.custom_vjp, nondiff_argnums=(0, 1))
+    def _clip_gradient(lo, hi, x):
+      return x  # identity function
+
+    def clip_gradient_fwd(lo, hi, x):
+        # return x, None
+        return x, (hi, )
+
+    def clip_gradient_bwd(lo, hi, _, g):
+        return (np.clip(g, lo, hi),)
+
+    _clip_gradient.defvjp(clip_gradient_fwd, clip_gradient_bwd)
+
+    def clip_gradient(x):
+        lo = -1
+        hi = x + 1  # causes things to break
+        return _clip_gradient(lo, hi, x)
+
+    jax.grad(clip_gradient)(1.)  # doesn't crash
+
 
 class DeprecatedCustomTransformsTest(jtu.JaxTestCase):