Repository: google/jax
Issue: Error when `custom_vjp`-wrapped original function closes over batch (vmap) tracer

Defining a `custom_vjp` and then calling its forward implementation (without even using `vjp` at all!) fails with a confusing error message when:
- the arguments to the `custom_vjp` function are not batched
- but the function closes over some value that is batched

```
def bad(x, y):
  @jax.custom_derivatives.custom_vjp
  def go(y): return x
  def never_called(*args): raise NotImplementedError()
  go.defvjp(never_called, never_called)
  return go(y)

bad(1., 2.)  # -> 1.
jax.vmap(bad)(jnp.array([1.]), jnp.array([2.]))  # -> [1.]
jax.vmap(bad, in_axes=[0, None])(jnp.array([1.]), 2.)  # -> error
```
The error in question:

```
<ipython-input-79-21aaf8a8397e> in <module>()
----> 1 jax.vmap(bad, in_axes=[0, None])(jnp.array([1.]), 2.)
```
<details>
<summary><code>--- (click to expand backtrace) ---</code></summary>
<pre><code>
.../jax/api.py in batched_fun(*args)
   1215                               lambda: flatten_axes("vmap out_axes", out_tree(),
   1216                                                    out_axes),
-&gt; 1217                               axis_name=axis_name)
   1218     return tree_unflatten(out_tree(), out_flat)
   1219 
<br><br>
.../jax/interpreters/batching.py in batch(fun, in_vals, in_dims, out_dim_dests, axis_name)
     33   # executes a batched version of `fun` following out_dim_dests
     34   batched_fun = batch_fun(fun, in_dims, out_dim_dests, axis_name=axis_name)
---&gt; 35   return batched_fun.call_wrapped(*in_vals)
     36 
     37 @lu.transformation_with_aux
<br><br>
.../jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    158 
    159     try:
--&gt; 160       ans = self.f(*args, **dict(self.params, **kwargs))
    161     except:
    162       # Some transformations yield from inside context managers, so we have to
<br><br>
&lt;ipython-input-75-0975376ba419&gt; in bad(x, y)
      4   def never_called(*args): raise NotImplementedError()
      5   go.defvjp(never_called, never_called)
----&gt; 6   return go(y)
<br><br>
.../jax/custom_derivatives.py in __call__(self, *args, **kwargs)
    488     if config.omnistaging_enabled:
    489       out_flat = custom_vjp_call_p.bind(flat_fun, flat_fwd, flat_bwd, *args_flat,
--&gt; 490                                         out_trees=out_trees)
    491       fst, aux = lu.merge_linear_aux(out_tree, out_trees)
    492       out_tree = aux if fst else aux[0]
<br><br>
.../jax/custom_derivatives.py in bind(self, fun, fwd, bwd, out_trees, *args)
    577     with core.maybe_new_sublevel(top_trace):
    578       outs = top_trace.process_custom_vjp_call(self, fun, fwd, bwd, tracers,
--&gt; 579                                                out_trees=out_trees)
    580     _, env_trace_todo = lu.merge_linear_aux(env_trace_todo1, env_trace_todo2)
    581     return _apply_todos(env_trace_todo, map(core.full_lower, outs))
<br><br>
.../jax/core.py in process_custom_vjp_call(***failed resolving arguments***)
    605   def process_custom_vjp_call(self, primitive, fun, fwd, bwd, tracers, out_trees):
    606     del primitive, fwd, bwd, out_trees  # Unused.
--&gt; 607     return fun.call_wrapped(*tracers)
    608 
    609 
<br><br>
.../jax/linear_util.py in call_wrapped(***failed resolving arguments***)
    171     while stack:
    172       gen, out_store = stack.pop()
--&gt; 173       ans = gen.send(ans)
    174       if out_store is not None:
    175         ans, side = ans
<br><br>
.../jax/core.py in process_env_traces(primitive, level, params_tuple, out_axes_transforms, *args)
   1195     trace = ans._trace.main.with_cur_sublevel()
   1196     outs = map(trace.full_raise, outs)
-&gt; 1197     outs, cur_todo = primitive.post_process(trace, outs, params)
   1198     if isinstance(primitive, MapPrimitive):
   1199       cur_todo, out_axes_transform = cur_todo
</code></pre>
</details>

```
.../jax/custom_derivatives.py in post_process(self, trace, out_tracers, params)
    586 
    587   def post_process(self, trace, out_tracers, params):
--> 588     return trace.post_process_custom_vjp_call(out_tracers, params)
    589 custom_vjp_call_p = CustomVJPCallPrimitive('custom_vjp_call')
    590 

AttributeError: 'BatchTrace' object has no attribute 'post_process_custom_vjp_call'
```

Note that the existing test `api_test.test_bwd_closes_over_tracer` test for closing over a tracer in the backward pass implementation. However, this error happens when closing over a tracer in the non-vjp implementation, and happens even if we don't ask for any gradients.

Changed files:
jax/interpreters/batching.py, tests/api_test.py

diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
index aefa6ee34..d738e5296 100644
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -259,6 +259,8 @@ class BatchTrace(Trace):
       out_dims = out_dims[-len(out_vals) % len(out_dims):]
     return [BatchTracer(self, v, d) for v, d in zip(out_vals, out_dims)]
 
+  post_process_custom_vjp_call = post_process_custom_jvp_call
+
 def _main_trace_for_axis_names(main_trace: core.MainTrace,
                                axis_name: Union[core.AxisName, Tuple[core.AxisName, ...]]
                                ) -> bool:
diff --git a/tests/api_test.py b/tests/api_test.py
index 20ef99463..4d5cc19cb 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -3866,6 +3866,47 @@ class CustomVJPTest(jtu.JaxTestCase):
     expected = jnp.cos(3.)
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def test_closed_over_tracer2(self):
+    def outer(x):
+      @api.custom_vjp
+      def f(y):
+        return x * y
+      def f_fwd(y):
+        return f(y), jnp.cos(y)
+      def f_rev(cos_y, g):
+        return (cos_y * g,)
+      f.defvjp(f_fwd, f_rev)
+      return f
+
+    @api.vmap
+    def g(x):
+      return outer(x)(3.)
+
+    ans = g(np.arange(3.))
+    expected = np.arange(3.) * 3
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+  def test_closed_over_tracer3(self):
+    def outer(x):
+      @api.custom_vjp
+      def f(y):
+        return x * y
+      def f_fwd(y):
+        return f(y), (x, jnp.cos(y))
+      def f_rev(res, g):
+        x, cos_y = res
+        return (cos_y * g * x,)
+      f.defvjp(f_fwd, f_rev)
+      return api.grad(f)
+
+    @api.vmap
+    def g(x):
+      return outer(x)(3.)
+
+    ans = g(np.arange(3.))
+    expected = np.cos(3.) * np.arange(3.)
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
   def test_nondiff_arg_tracer_error(self):
     # This is similar to the old (now skipped) test_nondiff_arg_tracer, except
     # we're testing for the error message that that usage pattern now raises.