Repository: google/jax
Issue: grad + vmap + odeint AssertionError

### Discussed in https://github.com/google/jax/discussions/8782

<div type='discussions-op-text'>

<sup>Originally posted by **DanPuzzuoli** December  2, 2021</sup>
Hi,

I've seen a bunch of discussions and issues surrounding this so I apologize if i'm re-raising something that has already been addressed elsewhere. I don't understand the internals of JAX enough to understand if this is some version of an issue that's already raised, though a lot of what I'm seeing is in already closed issues so I assume are solved and hence this is different.

When I run the following code on the latest release versions of jax/jaxlib (a self-contained version of my actual code):
```
from jax.experimental.ode import odeint
from jax import jit, value_and_grad, vmap
import jax.numpy as jnp
from jax.config import config
config.update("jax_enable_x64", True)
config.update('jax_platform_name', 'cpu')


T = 1.

X = -1j * jnp.array([[0., 1.], [1., 0.]], dtype=complex)
Y = -1j * jnp.array([[0., -1j], [1j, 0.]], dtype=complex)

def err_obj(a, b_vals):
    
    def err(b):
        def rhs(y, t):
            return (b * X + a * (t**2) * Y) @ y
        
        res = odeint(rhs, y0=jnp.eye(2, dtype=complex), t=jnp.array([0, T], dtype=complex), rtol=1e-6, atol=1e-6)
        
        return jnp.abs((X * res[-1]).sum())**2 / 4


    all_err = vmap(err)(b_vals)
    return all_err.sum()

b_vals = jnp.array([1., 2., 3., 4., 5.])
jit(value_and_grad(lambda a: err_obj(a, b_vals)))(1.)
``` 

I get the error:
```
AssertionError: length mismatch: [1, 4]
```

I've definitely reverse-mode differentiated this code in the past with success, though it was some time ago so would have been on a much older version of jax/jaxlib. Based on other similar issues I've seen it seems like this has something to do with the interaction of reverse-mode autodiff, vmap, and the control flow used in `odeint`, but again, the issues I've seen raising these kinds of errors seem to have been solved?

Thanks!</div>

Changed files:
jax/experimental/jax2tf/jax2tf.py, jax/interpreters/batching.py, jax/core.py, jax/interpreters/ad.py, jax/interpreters/partial_eval.py, tests/api_test.py, jax/_src/custom_derivatives.py, jax/experimental/maps.py

diff --git a/jax/_src/custom_derivatives.py b/jax/_src/custom_derivatives.py
index 923df3c7e..217e42bb3 100644
--- a/jax/_src/custom_derivatives.py
+++ b/jax/_src/custom_derivatives.py
@@ -274,10 +274,10 @@ class CustomJVPCallPrimitive(core.CallPrimitive):
   def bind(self, fun, jvp, *args):
     args = map(core.full_lower, args)
     top_trace = core.find_top_trace(args)
-    fun, env_trace_todo1 = core.process_env_traces(
-        fun, self, top_trace and top_trace.level, (), None)
-    jvp, env_trace_todo2 = core.process_env_traces(
-        jvp, self, top_trace and top_trace.level, (), None)
+    fun, env_trace_todo1 = process_env_traces(
+        fun, self, top_trace and top_trace.level, False)
+    jvp, env_trace_todo2 = process_env_traces(
+        jvp, self, top_trace and top_trace.level, True)
     tracers = map(top_trace.full_raise, args)  # type: ignore
     outs = top_trace.process_custom_jvp_call(self, fun, jvp, tracers)  # type: ignore
     _, env_trace_todo = lu.merge_linear_aux(env_trace_todo1, env_trace_todo2)
@@ -287,8 +287,25 @@ class CustomJVPCallPrimitive(core.CallPrimitive):
     with core.new_sublevel():
       return fun.call_wrapped(*args)
 
-  def post_process(self, trace, out_tracers, params):
-    return trace.post_process_custom_jvp_call(out_tracers, params)
+  def post_process(self, trace, out_tracers, jvp_was_run: bool):
+    return trace.post_process_custom_jvp_call(out_tracers, jvp_was_run)
+
+@lu.transformation_with_aux
+def process_env_traces(primitive, level: int, jvp_was_run: bool, *args):
+  outs = yield args, {}
+  todo = []
+  while True:
+    tracers = [x for x in outs if isinstance(x, core.Tracer)
+               and (level is None or x._trace.level > level)]
+    if tracers:
+      ans = max(tracers, key=lambda x: x._trace.level)
+    else:
+      break
+    trace = ans._trace.main.with_cur_sublevel()
+    outs = map(trace.full_raise, outs)
+    outs, cur_todo = primitive.post_process(trace, outs, jvp_was_run)
+    todo.append(cur_todo)
+  yield outs, tuple(todo)  # Ensure the aux output is immutable
 
 def _apply_todos(todos, outs):
   todos_list = list(todos)
@@ -567,6 +584,7 @@ def _flatten_fwd(in_tree, *args):
 @lu.transformation
 def _flatten_bwd(in_tree, in_avals, out_trees, *args):
   out_tree, res_tree = out_trees()
+  assert len(args) == res_tree.num_leaves + out_tree.num_leaves
   res, cts_out = split_list(args, [res_tree.num_leaves])
   py_res = tree_unflatten(res_tree, res)
   py_cts_out = tree_unflatten(out_tree, cts_out)
@@ -605,15 +623,21 @@ class CustomVJPCallPrimitive(core.CallPrimitive):
   def bind(self, fun, fwd, bwd, *args, out_trees):
     args = map(core.full_lower, args)
     top_trace = core.find_top_trace(args)
-    fun, env_trace_todo1 = core.process_env_traces(
-        fun, self, top_trace and top_trace.level, (), None)
-    fwd, env_trace_todo2 = core.process_env_traces(
-        fwd, self, top_trace and top_trace.level, (), None)
+    fun, env_trace_todo1 = process_env_traces(
+        fun, self, top_trace and top_trace.level, False)
+    fwd, env_trace_todo2 = process_env_traces_fwd(
+      fwd, top_trace and top_trace.level, out_trees)
     tracers = map(top_trace.full_raise, args)  # type: ignore
-    outs = top_trace.process_custom_vjp_call(self, fun, fwd, bwd, tracers,
-                                              out_trees=out_trees)
-    _, env_trace_todo = lu.merge_linear_aux(env_trace_todo1, env_trace_todo2)
-    return _apply_todos(env_trace_todo, map(core.full_lower, outs))
+    bwd_ = lu.wrap_init(lambda *args: bwd.call_wrapped(*args))
+    outs = top_trace.process_custom_vjp_call(self, fun, fwd, bwd_, tracers,
+                                             out_trees=out_trees)
+    fst, env_trace_todo = lu.merge_linear_aux(env_trace_todo1, env_trace_todo2)
+    if fst:
+      return _apply_todos(env_trace_todo, map(core.full_lower, outs))
+    else:
+      env_trace_todo, bwd_transform = env_trace_todo
+      bwd = _apply_bwd_transform(bwd_transform, bwd)
+      return _apply_todos(env_trace_todo, map(core.full_lower, outs))
 
   def impl(self, fun, fwd, bwd, *args, out_trees):
     del fwd, bwd, out_trees
@@ -624,6 +648,32 @@ class CustomVJPCallPrimitive(core.CallPrimitive):
     return trace.post_process_custom_vjp_call(out_tracers, params)
 custom_vjp_call_p = CustomVJPCallPrimitive('custom_vjp_call')
 
+@lu.transformation_with_aux
+def process_env_traces_fwd(level: int, out_trees, *args):
+  outs = yield args, {}
+  todo = []
+  bwd_transforms = []
+  while True:
+    tracers = [x for x in outs if isinstance(x, core.Tracer)
+               and (level is None or x._trace.level > level)]
+    if tracers:
+      ans = max(tracers, key=lambda x: x._trace.level)
+    else:
+      break
+    trace = ans._trace.main.with_cur_sublevel()
+    outs = map(trace.full_raise, outs)
+    outs, cur_todo, bwd_xform = trace.post_process_custom_vjp_call_fwd(outs, out_trees)
+    todo.append(cur_todo)
+    bwd_transforms.append(bwd_xform)
+  yield outs, (tuple(todo), tuple(bwd_transforms))
+
+
+def _apply_bwd_transform(todos, bwd):
+  todos_list = list(todos)
+  while todos_list:
+    bwd = todos_list.pop()(bwd)
+  return bwd
+
 def _custom_vjp_call_jaxpr_impl(*args, fun_jaxpr, **_):
   return core.jaxpr_as_fun(fun_jaxpr)(*args)
 
diff --git a/jax/core.py b/jax/core.py
index 3b5b0225d..db7fd96d7 100644
--- a/jax/core.py
+++ b/jax/core.py
@@ -1662,25 +1662,32 @@ def as_named_shape(shape) -> NamedShape:
 
 # ------------------- Call -------------------
 
-def apply_todos(todos, outs):
-  todos_list = list(todos)
-  while todos_list:
-    outs = map(full_lower, todos_list.pop()(outs))
-  return outs
+class CallPrimitive(Primitive):
+  multiple_results = True
+  call_primitive = True
 
-class _IgnoreElemList(list):
-  """Compares equal to all other _ignore_elem_lists."""
-  def __hash__(self): return 0
-  def __eq__(self, other):
-    return type(other) is _IgnoreElemList
+  def bind(self, fun, *args, **params):
+    return call_bind(self, fun, *args, **params)
+
+  def get_bind_params(self, params):
+    new_params = dict(params)
+    subfun = lu.wrap_init(partial(eval_jaxpr, new_params.pop('call_jaxpr'), ()))
+    return [subfun], new_params
+
+def call_bind(primitive: CallPrimitive, fun, *args, **params):
+  top_trace = find_top_trace(args)
+  fun, env_trace_todo = process_env_traces_call(
+      fun, primitive, top_trace and top_trace.level, tuple(params.items()))
+  tracers = map(top_trace.full_raise, args)
+  outs = top_trace.process_call(primitive, fun, tracers, params)
+  return map(full_lower, apply_todos(env_trace_todo(), outs))
 
 @lu.transformation_with_aux
-def process_env_traces(primitive: Union['CallPrimitive', 'MapPrimitive'],
-                       level: int, params_tuple: tuple, out_axes_transforms, *args):
+def process_env_traces_call(primitive: CallPrimitive, level: int,
+                            params_tuple: tuple, *args):
   outs = yield args, {}
   params = dict(params_tuple)
   todo = []
-  assert not out_axes_transforms
   while True:
     tracers = [x for x in outs if isinstance(x, Tracer)
                and (level is None or x._trace.level > level)]
@@ -1690,55 +1697,16 @@ def process_env_traces(primitive: Union['CallPrimitive', 'MapPrimitive'],
       break
     trace = ans._trace.main.with_cur_sublevel()
     outs = map(trace.full_raise, outs)
-    outs, cur_todo = primitive.post_process(trace, outs, params)
-    if isinstance(primitive, MapPrimitive):
-      cur_todo, out_axes_transform = cur_todo
-      out_axes_transforms.append(out_axes_transform)
+    outs, cur_todo = trace.post_process_call(primitive, outs, params)
     todo.append(cur_todo)
   yield outs, tuple(todo)  # Ensure the aux output is immutable
 
-def call_bind(primitive: Union['CallPrimitive', 'MapPrimitive'],
-              fun, *args, **params):
-  out_axes_transforms = _IgnoreElemList()
-  if primitive.map_primitive:
-    out_axes_thunk = params['out_axes_thunk']
-    # The new thunk depends deterministically on the old thunk and the wrapped
-    # function. Any caching already has to include the wrapped function as part
-    # of the key, so we only use the previous thunk for equality checks.
-    @as_hashable_function(closure=out_axes_thunk)
-    def new_out_axes_thunk():
-      out_axes = out_axes_thunk()
-      for t in out_axes_transforms:
-        out_axes = t(out_axes)
-      return out_axes
-    params = dict(params, out_axes_thunk=new_out_axes_thunk)
-  params_tuple = tuple(params.items())
-  top_trace = find_top_trace(args)
-  fun, env_trace_todo = process_env_traces(
-      fun, primitive, top_trace and top_trace.level,
-      params_tuple, out_axes_transforms)
-  tracers = map(top_trace.full_raise, args)
-  outs = primitive.process(top_trace, fun, tracers, params)
-  return map(full_lower, apply_todos(env_trace_todo(), outs))
-
-
-class CallPrimitive(Primitive):
-  multiple_results = True
-  call_primitive = True
-
-  def bind(self, fun, *args, **params):
-    return call_bind(self, fun, *args, **params)
-
-  def process(self, trace, fun, tracers, params):
-    return trace.process_call(self, fun, tracers, params)
+def apply_todos(todos, outs):
+  todos_list = list(todos)
+  while todos_list:
+    outs = map(full_lower, todos_list.pop()(outs))
+  return outs
 
-  def post_process(self, trace, out_tracers, params):
-    return trace.post_process_call(self, out_tracers, params)
-
-  def get_bind_params(self, params):
-    new_params = dict(params)
-    subfun = lu.wrap_init(partial(eval_jaxpr, new_params.pop('call_jaxpr'), ()))
-    return [subfun], new_params
 
 def call_impl(f: lu.WrappedFun, *args, **params):
   del params  # params parameterize the call primitive, not the function
@@ -1781,6 +1749,70 @@ def primitive_uses_outfeed(prim: Primitive, params: Dict) -> bool:
 
 # ------------------- Map -------------------
 
+class MapPrimitive(Primitive):
+  multiple_results = True
+  map_primitive = True
+
+  def bind(self, fun, *args, **params):
+    assert len(params['in_axes']) == len(args)
+    return map_bind(self, fun, *args, **params)
+
+  def process(self, trace, fun, tracers, params):
+    return trace.process_map(self, fun, tracers, params)
+
+  def post_process(self, trace, out_tracers, params):
+    return trace.post_process_map(self, out_tracers, params)
+
+  def get_bind_params(self, params):
+    new_params = dict(params)
+    subfun = lu.wrap_init(partial(eval_jaxpr, new_params.pop('call_jaxpr'), ()))
+    axes = new_params.pop('out_axes')
+    new_params['out_axes_thunk'] = HashableFunction(lambda: axes, closure=axes)
+    return [subfun], new_params
+
+def map_bind(primitive: 'MapPrimitive', fun, *args, out_axes_thunk, **params):
+  # The new thunk depends deterministically on the old thunk and the wrapped
+  # function. Any caching already has to include the wrapped function as part
+  # of the key, so we only use the previous thunk for equality checks.
+  @as_hashable_function(closure=out_axes_thunk)
+  def new_out_axes_thunk():
+    out_axes = out_axes_thunk()
+    _, out_axes_transforms = todo_and_xforms()
+    for t in out_axes_transforms:
+      out_axes = t(out_axes)
+    return out_axes
+  params = dict(params, out_axes_thunk=new_out_axes_thunk)
+  params_tuple = tuple(params.items())
+  top_trace = find_top_trace(args)
+  fun, todo_and_xforms = process_env_traces_map(
+      fun, primitive, top_trace and top_trace.level, params_tuple)
+  tracers = map(top_trace.full_raise, args)
+  outs = primitive.process(top_trace, fun, tracers, params)
+  env_trace_todo, _ = todo_and_xforms()
+  return map(full_lower, apply_todos(env_trace_todo, outs))
+
+@lu.transformation_with_aux
+def process_env_traces_map(primitive: MapPrimitive, level: int,
+                           params_tuple: tuple, *args):
+  outs = yield args, {}
+  params = dict(params_tuple)
+  todo = []
+  out_axes_transforms = []
+  while True:
+    tracers = [x for x in outs if isinstance(x, Tracer)
+               and (level is None or x._trace.level > level)]
+    if tracers:
+      ans = max(tracers, key=lambda x: x._trace.level)
+    else:
+      break
+    trace = ans._trace.main.with_cur_sublevel()
+    outs = map(trace.full_raise, outs)
+    outs, (cur_todo, cur_xform) = primitive.post_process(trace, outs, params)
+    todo.append(cur_todo)
+    out_axes_transforms.append(cur_xform)
+  yield outs, (tuple(todo), tuple(out_axes_transforms))
+
+
 def mapped_aval(size: int, axis: int, aval: AbstractValue) -> AbstractValue:
   handler, _ = aval_mapping_handlers.get(type(aval), (None, None))
   if handler is not None:
@@ -1818,28 +1850,6 @@ aval_mapping_handlers: Dict[Type, AvalMapHandlerPair] = {
     ConcreteArray: (_map_shaped_array, _unmap_shaped_array),
 }
 
-
-class MapPrimitive(Primitive):
-  multiple_results = True
-  map_primitive = True
-
-  def bind(self, fun, *args, **params):
-    assert len(params['in_axes']) == len(args)
-    return call_bind(self, fun, *args, **params)
-
-  def process(self, trace, fun, tracers, params):
-    return trace.process_map(self, fun, tracers, params)
-
-  def post_process(self, trace, out_tracers, params):
-    return trace.post_process_map(self, out_tracers, params)
-
-  def get_bind_params(self, params):
-    new_params = dict(params)
-    subfun = lu.wrap_init(partial(eval_jaxpr, new_params.pop('call_jaxpr'), ()))
-    axes = new_params.pop('out_axes')
-    new_params['out_axes_thunk'] = HashableFunction(lambda: axes, closure=axes)
-    return [subfun], new_params
-
 @contextmanager
 def extend_axis_env(axis_name: AxisName, size: int, tag: Any):
   frame = AxisEnvFrame(axis_name, size, tag)
diff --git a/jax/experimental/jax2tf/jax2tf.py b/jax/experimental/jax2tf/jax2tf.py
index c8ecde3f6..6cd6fc16d 100644
--- a/jax/experimental/jax2tf/jax2tf.py
+++ b/jax/experimental/jax2tf/jax2tf.py
@@ -928,7 +928,7 @@ class TensorFlowTrace(core.Trace):
     del jvp  # Unused.
     return self.process_call(core.call_p, fun, tracers, {})
 
-  def post_process_custom_jvp_call(self, out_tracers, params):
+  def post_process_custom_jvp_call(self, out_tracers, _):
     assert False  # unreachable assuming jax2tf runs with clean trace state
 
   def process_custom_vjp_call(self, prim, fun, fwd, bwd, tracers, out_trees):
@@ -938,7 +938,10 @@ class TensorFlowTrace(core.Trace):
     del fwd, bwd, out_trees  # Unused.
     return self.process_call(core.call_p, fun, tracers, {})
 
-  def post_process_custom_vjp_call(self, out_tracers, params):
+  def post_process_custom_vjp_call(self, out_tracers, _):
+    assert False  # unreachable assuming jax2tf runs with clean trace state
+
+  def post_process_custom_vjp_call_fwd(self, *_, **__):
     assert False  # unreachable assuming jax2tf runs with clean trace state
 
   def get_primitive_impl(self, p: core.Primitive) -> Tuple[Callable, bool]:
diff --git a/jax/experimental/maps.py b/jax/experimental/maps.py
index aa8ab776f..1e5f9d82a 100644
--- a/jax/experimental/maps.py
+++ b/jax/experimental/maps.py
@@ -850,15 +850,15 @@ class EvaluationPlan(NamedTuple):
 # -------- xmap primitive and its transforms --------
 
 # xmap has a different set of parameters than pmap, so we make it its own primitive type
-class XMapPrimitive(core.MapPrimitive):  # Not really a map, but it gives us a few good defaults
+class XMapPrimitive(core.MapPrimitive):
   def __init__(self):
     super().__init__('xmap')
     self.def_impl(xmap_impl)
     self.def_custom_bind(self.bind)
 
-  def bind(self, fun, *args, **params):
-    assert len(params['in_axes']) == len(args), (params['in_axes'], args)
-    return core.call_bind(self, fun, *args, **params)  # type: ignore
+  def bind(self, fun, *args, in_axes, **params):
+    assert len(in_axes) == len(args), (in_axes, args)
+    return core.map_bind(self, fun, *args, in_axes=in_axes, **params)
 
   def process(self, trace, fun, tracers, params):
     return trace.process_xmap(self, fun, tracers, params)
diff --git a/jax/interpreters/ad.py b/jax/interpreters/ad.py
index f9db2de49..1a6398aa4 100644
--- a/jax/interpreters/ad.py
+++ b/jax/interpreters/ad.py
@@ -359,7 +359,7 @@ class JVPTrace(Trace):
     tangents_out = map(recast_to_float0, primals_out, tangents_out)
     return map(partial(JVPTracer, self), primals_out, tangents_out)
 
-  def post_process_custom_jvp_call(self, out_tracers, params):
+  def post_process_custom_jvp_call(self, out_tracers, _):
     raise CustomJVPException()
 
   def process_custom_vjp_call(self, _, __, fwd, bwd, tracers, *, out_trees):
@@ -375,7 +375,7 @@ class JVPTrace(Trace):
     tangents_out = map(recast_to_float0, primals_out, tangents_out)
     return map(partial(JVPTracer, self), primals_out, tangents_out)
 
-  def post_process_custom_vjp_call(self, out_tracers, params):
+  def post_process_custom_vjp_call(self, out_tracers, _):
     raise CustomVJPException()
 
   def join(self, xt, yt):
diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
index 43af6b82a..726ed7167 100644
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -274,12 +274,17 @@ class BatchTrace(Trace):
       out_dims = out_dims[:len(out_dims) // 2]
     return [BatchTracer(self, v, d) for v, d in zip(out_vals, out_dims)]
 
-  def post_process_custom_jvp_call(self, out_tracers, params):
+  def post_process_custom_jvp_call(self, out_tracers, jvp_was_run):
     vals, dims = unzip2((t.val, t.batch_dim) for t in out_tracers)
     main = self.main
     def todo(vals):
       trace = main.with_cur_sublevel()
-      return map(partial(BatchTracer, trace), vals, dims)
+      if jvp_was_run:
+        primal_dims, tangent_dims = dims[:len(vals)], dims[len(vals):]
+        assert primal_dims == tangent_dims
+        return map(partial(BatchTracer, trace), vals, primal_dims)
+      else:
+        return map(partial(BatchTracer, trace), vals, dims)
     return vals, todo
 
   def process_custom_vjp_call(self, prim, fun, fwd, bwd, tracers, *, out_trees):
@@ -296,7 +301,29 @@ 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 post_process_custom_vjp_call(self, out_tracers, _):
+    vals, dims = unzip2((t.val, t.batch_dim) for t in out_tracers)
+    main = self.main
+    def todo(vals):
+      trace = main.with_cur_sublevel()
+      return map(partial(BatchTracer, trace), vals, dims)
+    return vals, todo
+
+  def post_process_custom_vjp_call_fwd(self, out_tracers, out_trees):
+    vals, dims = unzip2((t.val, t.batch_dim) for t in out_tracers)
+    axis_size, = {x.shape[d] for x, d in zip(vals, dims) if d is not not_mapped}
+    main, trace_type = self.main, self.main.trace_type
+    axis_name = self.axis_name
+    _, res_tree = out_trees()
+    num_res = res_tree.num_leaves
+    res_dims, primal_dims = split_list(dims, [num_res])
+    def todo(vals):
+      trace = main.with_cur_sublevel()
+      return map(partial(BatchTracer, trace), vals, primal_dims)
+    def bwd_transform(bwd):
+      return batch_custom_vjp_bwd(bwd, axis_name, axis_size, dims, (None, ),
+                                  trace_type)
+    return vals, todo, bwd_transform
 
 def _main_trace_for_axis_names(main_trace: core.MainTrace,
                                axis_name: Iterable[core.AxisName],
diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
index f048f70d3..3e35b5688 100644
--- a/jax/interpreters/partial_eval.py
+++ b/jax/interpreters/partial_eval.py
@@ -174,6 +174,7 @@ class JaxprTrace(Trace):
       return out_tracer
 
   # We use process_call to handle both call and map primitives.
+  # TODO(mattjj): split out process_call and process_map
   def process_call(self, primitive, f: lu.WrappedFun, tracers, params):
     if primitive in call_partial_eval_rules:
       return call_partial_eval_rules[primitive](self, primitive, f, tracers, params)
@@ -253,8 +254,8 @@ class JaxprTrace(Trace):
   process_map = process_call
 
   # We use post_process_call to handle both call and map primitives.
+  # TODO(mattjj): split out post_process_call and post_process_map
   def post_process_call(self, primitive, out_tracers, params):
-
     jaxpr, consts, env = tracers_to_jaxpr([], out_tracers)
     out_pvs, out_pv_consts = unzip2(t.pval for t in out_tracers)
     out = out_pv_consts + consts
@@ -262,19 +263,23 @@ class JaxprTrace(Trace):
     del consts, out_pv_consts
     main = self.main
 
-    if primitive.map_primitive:
-      out_axes = params['out_axes_thunk']()
-      sz = params['axis_size']
-      out_pvs = [None if pv is None else core.unmapped_aval(sz, params['axis_name'], ax, pv)
-                 for pv, ax in zip(out_pvs, out_axes)]
-
     def todo(x):
+      if primitive.map_primitive:
+        out_axes_ = params['out_axes_thunk']()
+        assert len(out_axes_) == len(out_pvs) + nconsts  # transformed below
+        out_axes = out_axes_[:len(out_pvs)]
+        sz = params['axis_size']
+        out_pvs_ = [core.unmapped_aval(sz, params['axis_name'], ax, pv)
+                    if pv is not None else None for pv, ax in zip(out_pvs, out_axes)]
+      else:
+        out_pvs_ = out_pvs
+
       n = len(jaxpr.outvars)
       out_pv_consts, consts = x[:n], x[n:]
       trace = JaxprTrace(main, core.cur_sublevel())
       const_tracers = map(trace.new_instantiated_const, consts)
       out_tracers = [JaxprTracer(trace, PartialVal((out_pv, out_pv_const)), None)
-                     for out_pv, out_pv_const in zip(out_pvs, out_pv_consts)]
+                     for out_pv, out_pv_const in zip(out_pvs_, out_pv_consts)]
       in_tracers = (*const_tracers, *map(trace.full_raise, env))
 
       new_params = dict(params, call_jaxpr=convert_constvars_jaxpr(jaxpr))
@@ -349,7 +354,7 @@ class JaxprTrace(Trace):
     for t in out_tracers: t.recipe = eqn
     return out_tracers
 
-  def post_process_custom_jvp_call(self, out_tracers, params):
+  def post_process_custom_jvp_call(self, out_tracers, _):
     # This path should only be reachable if we expose a partial eval API
     # unrelated to autodiff, since we raise an error when differentiation with
     # respect to values over which a custom_jvp function closes is detected.
@@ -390,7 +395,7 @@ class JaxprTrace(Trace):
     for t in out_tracers: t.recipe = eqn
     return out_tracers
 
-  def post_process_custom_vjp_call(self, out_tracers, params):
+  def post_process_custom_vjp_call(self, out_tracers, _):
     # This path should only be reachable if we expose a partial eval API
     # unrelated to autodiff, since we raise an error when differentiation with
     # respect to values over which a custom_vjp function closes is detected.
@@ -1435,7 +1440,7 @@ class DynamicJaxprTrace(core.Trace):
     self.frame.eqns.append(eqn)
     return out_tracers
 
-  def post_process_custom_jvp_call(self, out_tracers, params):
+  def post_process_custom_jvp_call(self, out_tracers, _):
     assert False  # unreachable
 
   def process_custom_vjp_call(self, prim, fun, fwd, bwd, tracers, out_trees):
@@ -1459,7 +1464,7 @@ class DynamicJaxprTrace(core.Trace):
     self.frame.eqns.append(eqn)
     return out_tracers
 
-  def post_process_custom_vjp_call(self, out_tracers, params):
+  def post_process_custom_vjp_call(self, out_tracers, _):
     assert False  # unreachable
 
 def _memoize(thunk):
diff --git a/tests/api_test.py b/tests/api_test.py
index d0c5613d3..3cb6400b2 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -4955,6 +4955,28 @@ class CustomJVPTest(jtu.JaxTestCase):
 
     api.vmap(sample)(jax.random.split(jax.random.PRNGKey(1), 3))  # don't crash
 
+  def test_closure_with_vmap2(self):
+    # https://github.com/google/jax/issues/8783
+    def h(z):
+      def f(x):
+        @jax.custom_jvp
+        def g(y):
+          return x * y
+
+        # NOTE: rule closes over vmap tracer
+        @g.defjvp
+        def g_jvp(primals, tangents):
+          (y,), (ydot,) = primals, tangents
+          return x * y, x * ydot
+
+        return g(z)  # NOTE: no vmapped arg
+
+      return jax.vmap(f)(jnp.arange(3., dtype='float32'))
+
+    primals, tangents = jax.jvp(h, (jnp.float32(1.),), (jnp.float32(2.),))
+    self.assertAllClose(primals ,     jnp.arange(3., dtype='float32'))
+    self.assertAllClose(tangents, 2 * jnp.arange(3., dtype='float32'))
+
   @unittest.skipIf(numpy_version == (1, 21, 0),
                    "https://github.com/numpy/numpy/issues/19305")
   def test_float0(self):
@@ -6199,6 +6221,27 @@ class CustomVJPTest(jtu.JaxTestCase):
     jtu.check_grads(batched_scan_over_mul, (x_batch, coeff), order=2,
                     modes=['rev'])
 
+  def test_closure_with_vmap2(self):
+    # https://github.com/google/jax/issues/8783
+    def h(z):
+      def f(x):
+        @jax.custom_vjp
+        def g(y):
+          return x * y
+
+        def g_fwd(y):
+          return x * y, (x, x * y, y)
+        def g_rev(res, w_bar):
+          x, *_ = res
+          return (x * w_bar,)
+        g.defvjp(g_fwd, g_rev)
+
+        return g(z)
+
+      return jax.vmap(f)(jnp.arange(3., dtype='float32')).sum()
+
+    jtu.check_grads(h, (jnp.float32(3.14),), order=1, modes=['rev'])
+
 
 def transpose_unary(f, x_example):
   def transposed(y):