Repository: jax-ml/jax
Pull request #28589 (2025-05-07 18:05:49)

Consolidate initial/final style custom_vjp primitives into one
Having separate primitives for initial and final style meant that we needed some duplicated logic so this should be a net win.

~~Most of the tests pass, but there are a couple of failures related to closed-over tracers.~~

All tests green after https://github.com/jax-ml/jax/pull/28605 goes in.


Patch:
diff --git a/jax/_src/checkify.py b/jax/_src/checkify.py
index 144cbaf5c..aa9bfe952 100644
--- a/jax/_src/checkify.py
+++ b/jax/_src/checkify.py
@@ -1079,17 +1079,17 @@ def lift_jvp(num_errs: int, num_consts: int,
     return [*primal_errs, *out_primals, *tangent_errs, *out_tangents]
   return lu.wrap_init(jvp, debug_info=jvp_jaxpr_fun.debug_info)
 
-def custom_vjp_call_jaxpr_rule(in_err, enabled_errors, *in_vals,
-                               fun_jaxpr: core.ClosedJaxpr,
-                               fwd_jaxpr_thunk, num_consts,
-                               bwd: lu.WrappedFun, out_trees,
-                               symbolic_zeros: bool):
+def custom_vjp_call_rule(in_err, enabled_errors, *in_vals,
+                         call_jaxpr: core.ClosedJaxpr,
+                         fwd_jaxpr_thunk, num_consts,
+                         bwd: lu.WrappedFun, out_trees,
+                         symbolic_zeros: bool):
   err_vals, err_tree = jtu.tree_flatten(in_err)
   num_errs = err_tree.num_leaves
   checkified_fun = lu.wrap_init(
-      functools.partial(checkify_jaxpr_flat, fun_jaxpr.jaxpr,
-                        fun_jaxpr.consts, enabled_errors, err_tree),
-      debug_info=fun_jaxpr.jaxpr.debug_info)
+      functools.partial(checkify_jaxpr_flat, call_jaxpr.jaxpr,
+                        call_jaxpr.consts, enabled_errors, err_tree),
+      debug_info=call_jaxpr.jaxpr.debug_info)
   checkified_fun, fun_metadata = _flatten_and_get_error_metadata_thunk(
       checkified_fun)
 
@@ -1097,13 +1097,13 @@ def custom_vjp_call_jaxpr_rule(in_err, enabled_errors, *in_vals,
     # TODO(lenamartens, sharadmv): why not checkify here?
     xs, zeros = args[::2], args[1::2]
     xs, zeros = xs[num_errs:], zeros[num_errs:]
-    fwd_jaxpr, fwd_consts = fwd_jaxpr_thunk(*zeros)
+    fwd_jaxpr, fwd_consts = fwd_jaxpr_thunk.call_wrapped(*zeros)
     xs_without_consts = xs[num_consts:]
     return core.eval_jaxpr(fwd_jaxpr, fwd_consts, *xs_without_consts)
 
   # TODO(necula): the fwd result_paths are not quite the same as fun_jaxpr
   checkified_fwd_wrapped = lu.wrap_init(checkified_fwd,
-                                        debug_info=fun_jaxpr.jaxpr.debug_info)
+                                        debug_info=fwd_jaxpr_thunk.debug_info)
   bwd_ = lu.wrap_init(lambda *args: (*(None,)*num_errs, *bwd.call_wrapped(*args)),
                       debug_info=bwd.debug_info)
   checkified_fwd_wrapped, fwd_out_tree = flatten_fun_output(checkified_fwd_wrapped)
@@ -1118,7 +1118,7 @@ def custom_vjp_call_jaxpr_rule(in_err, enabled_errors, *in_vals,
   else:
     out_err, out_vals = in_err, all_outs
   return out_err, out_vals
-error_checks[custom_derivatives.custom_vjp_call_jaxpr_p] = custom_vjp_call_jaxpr_rule
+error_checks[custom_derivatives.custom_vjp_call_p] = custom_vjp_call_rule
 
 
 def check_discharge_rule(error, enabled_errors, *args, err_tree, debug):
diff --git a/jax/_src/custom_derivatives.py b/jax/_src/custom_derivatives.py
index dc8fc90e3..dcd893f44 100644
--- a/jax/_src/custom_derivatives.py
+++ b/jax/_src/custom_derivatives.py
@@ -425,16 +425,14 @@ def _custom_jvp_call_typecheck(_, *in_avals, call_jaxpr, jvp_jaxpr_fun,
   return call_jaxpr.out_avals, call_jaxpr.effects
 core.custom_typechecks[custom_jvp_call_p] = _custom_jvp_call_typecheck
 
-def _custom_jvp_call_mlir_translation(ctx, *args, call_jaxpr, jvp_jaxpr_fun,
-                                      num_consts, symbolic_zeros):
-  del jvp_jaxpr_fun, num_consts, symbolic_zeros
+def _custom_jvp_vjp_call_lowering(ctx, *args, call_jaxpr, **_):
   consts = mlir._ir_consts(call_jaxpr.consts)
   out, tokens = mlir.jaxpr_subcomp(ctx.module_context, call_jaxpr.jaxpr,
                                    ctx.name_stack, ctx.tokens_in, consts,
                                    *args, dim_var_values=ctx.dim_var_values)
   ctx.set_tokens_out(tokens)
   return out
-mlir.register_lowering(custom_jvp_call_p, _custom_jvp_call_mlir_translation)
+mlir.register_lowering(custom_jvp_call_p, _custom_jvp_vjp_call_lowering)
 
 # If a (multi)linear function is defined with a custom jvp, then
 # custom_jvp_call_ can appear in jaxprs to be transposed. Since it's already
@@ -936,8 +934,8 @@ def _temporary_dtype_exception(a, a_) -> bool:
 def _temporary_shape_exception(a, a_) -> bool:
   return config.custom_vjp_disable_shape_check.value
 
-class CustomVJPCallPrimitive(core.CallPrimitive):
-  initial_style: core.Primitive
+class CustomVJPCallPrimitive(core.Primitive):
+  multiple_results = True
 
   def bind(self, *args, **params):
     return self._true_bind(*args, **params)
@@ -946,107 +944,70 @@ class CustomVJPCallPrimitive(core.CallPrimitive):
     fun, fwd, bwd, tracers = args[0], args[1], args[2], args[3:]
     return trace.process_custom_vjp_call(self, fun, fwd, bwd, tracers, **params)
 
-custom_vjp_call_p = CustomVJPCallPrimitive('custom_vjp_call')
+  def impl(self, fun, fwd, bwd, *args):
+    raise NotImplementedError
+
+  def get_bind_params(self, params):
+    new_params = dict(params)
+    call_jaxpr: core.ClosedJaxpr = new_params.pop('call_jaxpr')
+    num_consts: int = new_params.pop('num_consts')
+    fwd_jaxpr_thunk = new_params.pop('fwd_jaxpr_thunk')
+    fun = lu.wrap_init(core.jaxpr_as_fun(call_jaxpr),
+                       debug_info=call_jaxpr.jaxpr.debug_info)
+    fwd = lift_fwd(num_consts, fwd_jaxpr_thunk)
+    const_avals, _ = split_list(call_jaxpr.in_avals, [num_consts])
+    bwd = _handle_consts_in_bwd(new_params.pop('bwd'), const_avals)
+    return [fun, fwd, bwd], new_params
+
+def lift_fwd(num_consts: int, fwd_jaxpr_thunk: lu.WrappedFun) -> lu.WrappedFun:
+  def fwd(*args):
+    vals, zeros = args[::2], args[1::2]
+    assert len(vals) == len(zeros)
+    _, primals = split_list(vals, [num_consts])
+    const_zeros, in_zeros = split_list(zeros, [num_consts])
+    if any(const_zeros):
+      raise ad.CustomVJPException()
+    fwd_jaxpr, fwd_consts = fwd_jaxpr_thunk.call_wrapped(*in_zeros)
+    return core.eval_jaxpr(fwd_jaxpr, fwd_consts, *primals)
+  return lu.wrap_init(fwd, debug_info=fwd_jaxpr_thunk.debug_info)
 
-def _custom_vjp_call_jaxpr_impl(*args, fun_jaxpr, **_):
-  return core.jaxpr_as_fun(fun_jaxpr)(*args)
+@lu.transformation2
+def _handle_consts_in_bwd(f, const_avals, *args):
+  return [Zero(a) for a in const_avals] + list(f(*args))
 
-def _custom_vjp_call_jaxpr_abstract_eval(*_, fun_jaxpr, **__):
-  disallowed_effects = effects.custom_derivatives_allowed_effects.filter_not_in(fun_jaxpr.effects)
+custom_vjp_call_p = CustomVJPCallPrimitive('custom_vjp_call')
+mlir.register_lowering(custom_vjp_call_p, _custom_jvp_vjp_call_lowering)
+
+def _custom_vjp_call_typecheck(_, *in_avals, call_jaxpr, **kwargs):
+  del in_avals, kwargs
+  disallowed_effects = effects.custom_derivatives_allowed_effects.filter_not_in(
+      call_jaxpr.effects)
   if disallowed_effects:
     raise NotImplementedError(
         f'Effects not supported in `custom_vjp`: {disallowed_effects}')
-  return fun_jaxpr.out_avals, fun_jaxpr.effects
-
-custom_vjp_call_jaxpr_p = core.Primitive('custom_vjp_call_jaxpr')
-custom_vjp_call_jaxpr_p.multiple_results = True
-custom_vjp_call_jaxpr_p.def_impl(_custom_vjp_call_jaxpr_impl)
-custom_vjp_call_jaxpr_p.def_effectful_abstract_eval(_custom_vjp_call_jaxpr_abstract_eval)
-CustomVJPCallPrimitive.initial_style = custom_vjp_call_jaxpr_p
-
-mlir.register_lowering(custom_vjp_call_jaxpr_p, mlir.lower_fun(
-    _custom_vjp_call_jaxpr_impl, multiple_results=True))
-
-def _custom_vjp_call_jaxpr_jvp(
-    primals, tangents, *, fun_jaxpr: core.ClosedJaxpr,
-    fwd_jaxpr_thunk: Callable[..., tuple[core.Jaxpr, Sequence[Any]]],
-    num_consts: int, bwd: lu.WrappedFun,
-    out_trees: Callable[[], Sequence[PyTreeDef]],
-    symbolic_zeros: bool):
-  _, args = split_list(primals, [num_consts])
-  consts_dot, args_dot = split_list(tangents, [num_consts])
-  if any(type(t) is not Zero for t in consts_dot):
-    raise ad.CustomVJPException()
-  zeros = [type(t) is not Zero for t in args_dot]
-  fwd_jaxpr, fwd_consts = fwd_jaxpr_thunk(*zeros)  # consts can be tracers!
-  _, res_tree = out_trees()
-  res_and_primals_out = core.eval_jaxpr(fwd_jaxpr, fwd_consts, *args)
-  res, primals_out = split_list(res_and_primals_out, [res_tree.num_leaves])
-  avals_out = [core.get_aval(x).to_tangent_aval() for x in primals_out]
-  args_dot = map(ad.instantiate_zeros, args_dot)
-  tangents_out = ad.custom_lin_p.bind(
-      *res, *args_dot, num_res=res_tree.num_leaves, bwd=bwd,
-      out_avals=avals_out, symbolic_zeros=symbolic_zeros)
-  tangents_out = map(lax.tie_p.bind, primals_out, tangents_out)
-  return primals_out, tangents_out
-ad.primitive_jvps[custom_vjp_call_jaxpr_p] = _custom_vjp_call_jaxpr_jvp
-
-def _custom_vjp_call_jaxpr_vmap(
-    axis_data, args, in_dims, *,
-    fun_jaxpr: core.ClosedJaxpr,
-    fwd_jaxpr_thunk: Callable[..., tuple[core.Jaxpr, Sequence[Any]]],
-    num_consts: int, bwd: lu.WrappedFun,
-    out_trees: Callable, symbolic_zeros: bool):
-  args = [batching.moveaxis(x, d, 0) if d is not not_mapped and d != 0
-          else x for x, d in zip(args, in_dims)]
-  in_batched = [d is not not_mapped for d in in_dims]
-  _, args_batched = split_list(in_batched, [num_consts])
-  batched_fun_jaxpr, out_batched = batching.batch_jaxpr(
-      fun_jaxpr, axis_data, in_batched, False)
-  out_dims1 = [0 if b else not_mapped for b in out_batched]
-  out_dims2 = []
-
-  @pe._memoize
-  def batched_fwd_jaxpr_thunk(*zeros):
-    fwd_jaxpr = core.ClosedJaxpr(*fwd_jaxpr_thunk(*zeros))  # consts can be tracers
-    batched_fwd_jaxpr, out_batched = batching.batch_jaxpr(
-        fwd_jaxpr, axis_data, args_batched, False)
-    out_dims2.append([0 if b else not_mapped for b in out_batched])
-    return batched_fwd_jaxpr.jaxpr, batched_fwd_jaxpr.consts
-
-  fwd_args_batched = [0 if b else not_mapped for b in args_batched]
-  fwd_out_dims = lambda: out_dims2[0]
-  tag = core.TraceTag()
-  batched_bwd = batching.batch_custom_vjp_bwd(
-    bwd, tag, axis_data, fwd_out_dims, fwd_args_batched)
-
-  batched_outs = custom_vjp_call_jaxpr_p.bind(
-      *args, fun_jaxpr=batched_fun_jaxpr,
-      fwd_jaxpr_thunk=batched_fwd_jaxpr_thunk, bwd=batched_bwd,
-      num_consts=num_consts, out_trees=out_trees, symbolic_zeros=symbolic_zeros)
-  out_dims = out_dims2[0] if out_dims2 else out_dims1
-  return batched_outs, out_dims
-batching.fancy_primitive_batchers[custom_vjp_call_jaxpr_p] = _custom_vjp_call_jaxpr_vmap
+  return call_jaxpr.out_avals, call_jaxpr.effects
+core.custom_typechecks[custom_vjp_call_p] = _custom_vjp_call_typecheck
 
-def _custom_vjp_call_jaxpr_dce(
+def _custom_vjp_call_dce(
     used_outs: Sequence[bool], eqn: core.JaxprEqn
 ) -> tuple[list[bool], core.JaxprEqn | None]:
   if not any(used_outs) and not pe.has_effects(eqn):
     return [False] * len(eqn.invars), None
-  fun_jaxpr: core.ClosedJaxpr = eqn.params["fun_jaxpr"]
+  call_jaxpr: core.ClosedJaxpr = eqn.params["call_jaxpr"]
   fwd_jaxpr_thunk = eqn.params["fwd_jaxpr_thunk"]
   bwd: lu.WrappedFun = eqn.params["bwd"]
   out_trees: Callable[[], Sequence[PyTreeDef]] = eqn.params["out_trees"]
   symbolic_zeros: bool = eqn.params["symbolic_zeros"]
-  dce_fun_jaxpr: core.ClosedJaxpr
+  dce_call_jaxpr: core.ClosedJaxpr
   used_ins: Sequence[bool]
-  dce_fun_jaxpr, used_ins = _cached_closed_call_dce_instantiate(
-      fun_jaxpr, tuple(used_outs))
+  dce_call_jaxpr, used_ins = _cached_closed_call_dce_instantiate(
+      call_jaxpr, tuple(used_outs))
   assert all(used_ins)
 
+  @partial(lu.wrap_init, debug_info=fwd_jaxpr_thunk.debug_info)
   @pe._memoize
   def dce_fwd_jaxpr_thunk(*zeros):
-    fwd_jaxpr = core.ClosedJaxpr(*fwd_jaxpr_thunk(*zeros))
+    fwd_jaxpr = core.ClosedJaxpr(*fwd_jaxpr_thunk.call_wrapped(*zeros))
     _, res_tree = out_trees()
     num_res = res_tree.num_leaves
     dce_fwd_jaxpr, _ = _cached_closed_call_dce_instantiate(
@@ -1058,7 +1019,7 @@ def _custom_vjp_call_jaxpr_dce(
     res, cts = split_list(args, [res_tree.num_leaves])
     cts_ = iter(cts)
     all_cts = []
-    for used, aval in zip(used_outs, fun_jaxpr.out_avals):
+    for used, aval in zip(used_outs, call_jaxpr.out_avals):
       if used:
         all_cts.append(next(cts_))
       else:
@@ -1075,17 +1036,15 @@ def _custom_vjp_call_jaxpr_dce(
   outvars = [v for used, v in zip(used_outs, eqn.outvars) if used]
   new_params = dict(
       eqn.params,
-      fun_jaxpr=dce_fun_jaxpr,
+      call_jaxpr=dce_call_jaxpr,
       fwd_jaxpr_thunk=dce_fwd_jaxpr_thunk,
       bwd=dce_bwd_wrapped,
   )
   new_eqn = pe.new_jaxpr_eqn(
-      eqn.invars, outvars, eqn.primitive, new_params, dce_fun_jaxpr.effects,
+      eqn.invars, outvars, eqn.primitive, new_params, dce_call_jaxpr.effects,
       eqn.source_info, eqn.ctx)
   return list(used_ins), new_eqn
-pe.dce_rules[custom_vjp_call_jaxpr_p] = _custom_vjp_call_jaxpr_dce
-
-xla.register_initial_style_primitive(custom_vjp_call_jaxpr_p)
+pe.dce_rules[custom_vjp_call_p] = _custom_vjp_call_dce
 
 batching.primitive_batchers[ad.custom_lin_p] = ad.raise_custom_vjp_error_on_jvp
 mlir.register_lowering(ad.custom_lin_p, ad.raise_custom_vjp_error_on_jvp)
@@ -1586,7 +1545,6 @@ def custom_vjp_by_custom_transpose(fun, fwd, bwd):
 # TODO(mattjj): remove these stubs, which exist to avoid breaking internal users
 custom_jvp_call_jaxpr_p = core.Primitive("custom_jvp_call_jaxpr")
 
-
 # The following is a helper for optimizing the behavior of custom_vjp when used
 # under remat. This is really only useful when the `fwd` function to custom_vjp
 # executes a black box kernel. Otherwise, DCE will perform this optimization
diff --git a/jax/_src/interpreters/partial_eval.py b/jax/_src/interpreters/partial_eval.py
index 64226a789..5866b0c5f 100644
--- a/jax/_src/interpreters/partial_eval.py
+++ b/jax/_src/interpreters/partial_eval.py
@@ -434,49 +434,45 @@ class JaxprTrace(Trace['JaxprTracer']):
     if all(t.is_known() for t in tracers):
       vals = [t.pval[1] for t in tracers]
       with core.set_current_trace(self.parent_trace):
-        return prim.bind(f, fwd, bwd, *vals, out_trees=out_trees, symbolic_zeros=symbolic_zeros)
-    else:
-      # TODO(mattjj): remove non-ad users of partial eval, then drop this case.
-      # We stage out the whole thing, i.e. no nontrivial partial evaluation.
-      tracers = map(self.instantiate_const_abstracted, tracers)
-      # Because we instantiate all tracers, in_knowns is all False.
-      in_knowns, in_avals, () = partition_pvals([t.pval for t in tracers])
-      f = trace_to_subjaxpr_nounits(f, self, True, f.debug_info)
-      f, aux = partial_eval_wrapper_nounits(f, (*in_knowns,), (*in_avals,))
-      with core.set_current_trace(self.parent_trace):
-        out_flat = prim.bind(f, fwd, bwd, out_trees=out_trees,
-                             symbolic_zeros=symbolic_zeros)
-      out_knowns, out_avals, jaxpr, env = aux()
-      out_consts, res = split_list(out_flat, [len(out_flat)-len(jaxpr.constvars)])
-      res_tracers = map(self.new_instantiated_const, res)
-      env_tracers = map(self.to_jaxpr_tracer, env)
-      out_tracers = [JaxprTracer(self, PartialVal.unknown(a), None)
-                    for a in out_avals]
-      closed_jaxpr = core.ClosedJaxpr(convert_constvars_jaxpr(jaxpr), ())
-
-      @_memoize
-      def fwd_jaxpr_thunk(*zeros):
-        fwd_ = _interleave_fun(fwd, zeros)
-        fwd_ = trace_to_subjaxpr_nounits(fwd_, self, True, fwd_.debug_info)
-        fwd_, aux = partial_eval_wrapper_nounits(fwd_, (*in_knowns,), (*in_avals,))
-        out_flat = fwd_.call_wrapped()
-        out_knowns, out_avals, jaxpr, env = aux()
-        _, res = split_list(out_flat, [len(out_flat)-len(jaxpr.constvars)])
-        converted_jaxpr = convert_envvars_to_constvars(jaxpr, len(env))
-        return converted_jaxpr, (*res, *env)
+        return prim.bind(f, fwd, bwd, *vals, out_trees=out_trees,
+                         symbolic_zeros=symbolic_zeros)
+
+    tracers = map(self.instantiate_const, tracers)
+    in_knowns = (False,) * len(tracers)
+    in_avals = tuple(t.aval for t in tracers)
+    f_ = trace_to_subjaxpr_nounits2(f, self.tag, f.debug_info, True)
+    f_, aux = partial_eval_wrapper_nounits(f_, in_knowns, in_avals)
+    params = dict(out_trees=out_trees, symbolic_zeros=symbolic_zeros)
+    res = prim.bind_with_trace(self.parent_trace, (f_, fwd, bwd), params)
+    out_knowns, out_avals, jaxpr, env = aux()
+    assert not any(out_knowns)
+    res_tracers = map(self.instantiate_const, map(self.new_const, res))
+    env_tracers = map(self.to_jaxpr_tracer, env)
+    out_tracers = [JaxprTracer(self, PartialVal.unknown(a), None)
+                   for a in out_avals]
+    closed_jaxpr = close_jaxpr(convert_constvars_jaxpr(jaxpr))
+
+    @partial(lu.wrap_init, debug_info=fwd.debug_info)
+    @_memoize
+    def fwd_jaxpr_thunk(*zeros):
+      fwd_ = _interleave_fun(fwd, zeros)
+      fwd_jaxpr, _, consts, () = trace_to_jaxpr_dynamic(fwd_, in_avals)
+      return fwd_jaxpr, consts
 
     name_stack = self._current_truncated_name_stack()
     source = source_info_util.current().replace(name_stack=name_stack)
+    params = dict(
+        call_jaxpr=closed_jaxpr,
+        fwd_jaxpr_thunk=fwd_jaxpr_thunk,
+        num_consts=len(res) + len(env),
+        bwd=bwd,
+        out_trees=out_trees,
+        symbolic_zeros=symbolic_zeros
+    )
     eqn = new_eqn_recipe((*res_tracers, *env_tracers, *tracers),
-                         out_tracers, prim.initial_style,
-                         dict(fun_jaxpr=closed_jaxpr,
-                              fwd_jaxpr_thunk=fwd_jaxpr_thunk,
-                              num_consts=len(res) + len(env),
-                              bwd=bwd, out_trees=out_trees,
-                              symbolic_zeros=symbolic_zeros),
-                         jaxpr.effects, source)
+                         out_tracers, prim, params, jaxpr.effects, source)
     for t in out_tracers: t.recipe = eqn
-    return merge_lists(out_knowns, out_tracers, out_consts)
+    return out_tracers
 
 def partition_pvals(
     pvals: list[PartialVal]
@@ -2050,6 +2046,7 @@ class DynamicJaxprTrace(core.Trace):
     fun_jaxpr, out_avals, consts, () = trace_to_jaxpr_dynamic(fun, in_avals)
     closed_fun_jaxpr = core.ClosedJaxpr(convert_constvars_jaxpr(fun_jaxpr), ())
 
+    @partial(lu.wrap_init, debug_info=jvp.debug_info)
     @_memoize
     def jvp_jaxpr_thunk(*in_zeros):
       for store in jvp.stores: store and store.reset()
@@ -2065,8 +2062,7 @@ class DynamicJaxprTrace(core.Trace):
     outvars = map(self.makevar, out_tracers)
     eqn = new_jaxpr_eqn([*constvars, *invars], outvars, prim,
                         dict(call_jaxpr=closed_fun_jaxpr,
-                             jvp_jaxpr_fun=lu.wrap_init(jvp_jaxpr_thunk,
-                                                        debug_info=jvp.debug_info),
+                             jvp_jaxpr_fun=jvp_jaxpr_thunk,
                              num_consts=len(consts),
                              symbolic_zeros=symbolic_zeros),
                         fun_jaxpr.effects,
@@ -2086,6 +2082,7 @@ class DynamicJaxprTrace(core.Trace):
     fun_jaxpr, out_avals, consts, _ = trace_to_jaxpr_dynamic(fun, in_avals)
     closed_fun_jaxpr = core.ClosedJaxpr(convert_constvars_jaxpr(fun_jaxpr), ())
 
+    @partial(lu.wrap_init, debug_info=fwd.debug_info)
     @_memoize
     def fwd_jaxpr_from_zeros(*zeros):
       for store in fwd.stores: store and store.reset()
@@ -2098,9 +2095,8 @@ class DynamicJaxprTrace(core.Trace):
     invars = map(self.getvar, tracers)
     constvars = map(self.getvar, map(to_jaxpr_tracer, consts))
     outvars = map(self.makevar, out_tracers)
-    eqn = new_jaxpr_eqn([*constvars, *invars], outvars,
-                        prim.initial_style,  # pytype: disable=attribute-error
-                        dict(fun_jaxpr=closed_fun_jaxpr,
+    eqn = new_jaxpr_eqn([*constvars, *invars], outvars, prim,
+                        dict(call_jaxpr=closed_fun_jaxpr,
                              fwd_jaxpr_thunk=fwd_jaxpr_from_zeros,
                              num_consts=len(consts),
                              bwd=bwd, out_trees=out_trees,
diff --git a/jax/_src/pallas/cost_estimate.py b/jax/_src/pallas/cost_estimate.py
index 3b82d3095..93bcf5348 100644
--- a/jax/_src/pallas/cost_estimate.py
+++ b/jax/_src/pallas/cost_estimate.py
@@ -238,15 +238,15 @@ def _pjit_cost_rule(ctx, *, jaxpr: jax_core.ClosedJaxpr, **_):
   )
 register_cost_rule(pjit.pjit_p, _pjit_cost_rule)
 
-def _custom_vjp_rule(ctx, *, fun_jaxpr: jax_core.ClosedJaxpr, **_):
+def _custom_vjp_rule(ctx, *, call_jaxpr: jax_core.ClosedJaxpr, **_):
   del ctx
-  inner_cost = cost_estimate_jaxpr(fun_jaxpr)
+  inner_cost = cost_estimate_jaxpr(call_jaxpr)
   return CostEstimate(
       flops=inner_cost.flops,
       transcendentals=inner_cost.transcendentals,
       bytes_accessed=inner_cost.bytes_accessed,
   )
-register_cost_rule(custom_derivatives.custom_vjp_call_jaxpr_p, _custom_vjp_rule)
+register_cost_rule(custom_derivatives.custom_vjp_call_p, _custom_vjp_rule)
 
 def _run_state_rule(*_, jaxpr: jax_core.Jaxpr, **_2):
   inner_cost = cost_estimate_jaxpr(pe.close_jaxpr(jaxpr))
diff --git a/jax/custom_derivatives.py b/jax/custom_derivatives.py
index 3628ae4aa..b768b687d 100644
--- a/jax/custom_derivatives.py
+++ b/jax/custom_derivatives.py
@@ -26,7 +26,6 @@ from jax._src.custom_derivatives import (
   custom_jvp_call_jaxpr_p as custom_jvp_call_jaxpr_p,
   custom_vjp as custom_vjp,
   custom_vjp_call_p as custom_vjp_call_p,
-  custom_vjp_call_jaxpr_p as custom_vjp_call_jaxpr_p,
   custom_vjp_primal_tree_values as custom_vjp_primal_tree_values,
   CustomVJPPrimal as CustomVJPPrimal,
   linear_call as linear_call,
diff --git a/jax/experimental/jax2tf/jax2tf.py b/jax/experimental/jax2tf/jax2tf.py
index 786e021e2..536bf1f20 100644
--- a/jax/experimental/jax2tf/jax2tf.py
+++ b/jax/experimental/jax2tf/jax2tf.py
@@ -3461,14 +3461,14 @@ def _custom_jvp_call(*args: TfVal, call_jaxpr: core.ClosedJaxpr,
 tf_impl[custom_derivatives.custom_jvp_call_p] = _custom_jvp_call
 
 
-def _custom_vjp_call_jaxpr(*args: TfVal, fun_jaxpr: core.ClosedJaxpr,
-                           **_) -> Sequence[TfVal]:
+def _custom_vjp_call(*args: TfVal, call_jaxpr: core.ClosedJaxpr,
+                     **_) -> Sequence[TfVal]:
   # TODO(necula): ensure that there is no AD transformation in scope
-  return _interpret_jaxpr(fun_jaxpr, *args, extra_name_stack="custom_vjp",
+  return _interpret_jaxpr(call_jaxpr, *args, extra_name_stack="custom_vjp",
                           fresh_constant_cache=False)
 
 
-tf_impl[custom_derivatives.custom_vjp_call_jaxpr_p] = _custom_vjp_call_jaxpr
+tf_impl[custom_derivatives.custom_vjp_call_p] = _custom_vjp_call
 
 
 def _custom_lin(*args: TfVal, **_) -> Sequence[TfVal]:
diff --git a/jax/extend/core/primitives.py b/jax/extend/core/primitives.py
index 60d8cd24a..30350dace 100644
--- a/jax/extend/core/primitives.py
+++ b/jax/extend/core/primitives.py
@@ -26,7 +26,6 @@ from jax._src.custom_derivatives import (
   custom_jvp_call_p as custom_jvp_call_p,
   custom_jvp_call_jaxpr_p as custom_jvp_call_jaxpr_p,
   custom_vjp_call_p as custom_vjp_call_p,
-  custom_vjp_call_jaxpr_p as custom_vjp_call_jaxpr_p,
 )
 
 from jax._src.dispatch import device_put_p as device_put_p


Test patch:
diff --git a/tests/custom_api_test.py b/tests/custom_api_test.py
index 72c14634a..73dc2fbef 100644
--- a/tests/custom_api_test.py
+++ b/tests/custom_api_test.py
@@ -3065,7 +3065,7 @@ class CustomVJPTest(jtu.JaxTestCase):
       if not dce_jaxpr.eqns:
         assert not includes
         return
-      call_jaxpr = dce_jaxpr.eqns[0].params["fun_jaxpr"]
+      call_jaxpr = dce_jaxpr.eqns[0].params["call_jaxpr"]
       for prim in includes:
         assert any(eqn.primitive == prim for eqn in call_jaxpr.eqns)
       for prim in excludes: