# Copyright 2020 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from collections.abc import Callable, Sequence
import dataclasses
from functools import update_wrapper, reduce, partial, wraps
from typing import Any

from jax._src import config
from jax._src import core
from jax._src import custom_api_util
from jax._src import dtypes
from jax._src import effects
from jax._src import flattree as ft
from jax._src import linear_util as lu
from jax._src import traceback_util
from jax._src.ad_util import (
    stop_gradient_p, SymbolicZero, Zero, zeros_like_aval, p2tz)
from jax._src.api_util import (
  argnums_partial, resolve_kwargs,
  prepend_static_args, debug_info, fun_signature,
  infer_argnums_and_argnames, fun_sourceinfo)
from jax._src.errors import UnexpectedTracerError
from jax._src.state.types import AbstractRef
from jax._src.interpreters import ad
from jax._src.interpreters import batching
from jax._src.interpreters import mlir
from jax._src.interpreters import partial_eval as pe
from jax._src.tree_util import (
    tree_flatten, tree_unflatten, tree_map, treedef_is_leaf, treedef_tuple,
    register_pytree_node_class, tree_leaves, tree_flatten_with_path,
    tree_leaves_with_path, keystr, treedef_children, tree_structure, PyTreeDef,
    tracing_registry, broadcast_prefix)
from jax._src.util import (cache, safe_zip, safe_map, split_list, unzip2,
                           weakref_lru_cache)


traceback_util.register_exclusion(__file__)

map = safe_map
zip = safe_zip


### util

def _sum_tangents(_, x, *xs):
  return reduce(ad.add_tangents, xs, x)

def _zeros_like_pytree(x):
  return tree_map(p2tz, x)

_stop_gradient = partial(
    tree_map,
    lambda x: stop_gradient_p.bind(x) if isinstance(x, core.Tracer) else x,
)


# like the api_util.py function, but also grabs output avals for error checking
@lu.transformation_with_aux2
def _flatten_fun_nokwargs(f: Callable,
                          store: lu.Store, in_tree: PyTreeDef,
                          *args_flat):
  py_args = tree_unflatten(in_tree, args_flat)
  ans = f(*py_args)
  ans_flat, ans_tree = tree_flatten(ans)
  ans_avals = [core.typeof(x) for x in ans_flat]
  store.store((ans_tree, ans_avals, ()))
  return ans_flat


### JVPs

@custom_api_util.register_custom_decorator_type
class custom_jvp[ReturnValue]:
  """Set up a JAX-transformable function for a custom JVP rule definition.

  This class is meant to be used as a function decorator. Instances are
  callables that behave similarly to the underlying function to which the
  decorator was applied, except when a differentiation transformation (like
  :py:func:`jax.jvp` or :py:func:`jax.grad`) is applied, in which case a custom
  user-supplied JVP rule function is used instead of tracing into and
  performing automatic differentiation of the underlying function's
  implementation.

  There are two instance methods available for defining the custom JVP rule:
  :py:func:`~jax.custom_jvp.defjvp` for defining a *single* custom JVP rule for
  all the function's inputs, and for convenience
  :py:func:`~jax.custom_jvp.defjvps`, which wraps
  :py:func:`~jax.custom_jvp.defjvp`, and allows you to provide separate
  definitions for the partial derivatives of the function w.r.t. each of its
  arguments.

  For example::

    @jax.custom_jvp
    def f(x, y):
      return jnp.sin(x) * y

    @f.defjvp
    def f_jvp(primals, tangents):
      x, y = primals
      x_dot, y_dot = tangents
      primal_out = f(x, y)
      tangent_out = jnp.cos(x) * x_dot * y + jnp.sin(x) * y_dot
      return primal_out, tangent_out

  For a more detailed introduction, see the tutorial_.

  .. _tutorial: https://docs.jax.dev/en/latest/301/custom-jvp-vjp.html
  """
  fun: Callable[..., ReturnValue]
  nondiff_argnums: Sequence[int]
  nondiff_argnames: Sequence[str]
  jvp: Callable[..., tuple[ReturnValue, ReturnValue]] | None = None
  symbolic_zeros: bool = False

  def __init__(self,
               fun: Callable[..., ReturnValue],
               nondiff_argnums: Sequence[int] = (),
               nondiff_argnames: Sequence[str] = (),
               ):
    update_wrapper(self, fun)
    self.fun = fun

    nondiff_argnums_: set[int] = set()
    if nondiff_argnames:
      sig = fun_signature(self.fun)
      assert sig is not None
      inferred_nondiff_argnums, _ = infer_argnums_and_argnames(
          sig, None, nondiff_argnames
      )
      nondiff_argnums_.update(inferred_nondiff_argnums)

    if nondiff_argnums:
      nondiff_argnums_.update(nondiff_argnums)

    self.nondiff_argnums = tuple(sorted(nondiff_argnums_))

  __getattr__ = custom_api_util.forward_attr

  def defjvp(self,
             jvp: Callable[..., tuple[ReturnValue, ReturnValue]],
             symbolic_zeros: bool = False,
             ) -> Callable[..., tuple[ReturnValue, ReturnValue]]:
    """Define a custom JVP rule for the function represented by this instance.

    Args:
      jvp: a Python callable representing the custom JVP rule. When there are no
        ``nondiff_argnums``, the ``jvp`` function should accept two arguments,
        where the first is a tuple of primal inputs and the second is a tuple of
        tangent inputs. The lengths of both tuples are equal to the number of
        parameters of the :class:`~jax.custom_jvp` function. The ``jvp`` function
        should produce as output a pair where the first element is the primal
        output and the second element is the tangent output. Elements of the
        input and output tuples may be arrays or any nested tuples/lists/dicts
        thereof.
      symbolic_zeros: boolean, indicating whether the rule should be passed
        objects representing static symbolic zeros in its tangent argument in
        correspondence with unperturbed values; otherwise, only standard JAX
        types (e.g. array-likes) are passed. Setting this option to ``True``
        allows a JVP rule to detect whether certain inputs are not involved in
        differentiation, but at the cost of needing special handling for these
        objects (which e.g. can't be passed into jax.numpy functions). Default
        ``False``.

    Returns:
      Returns ``jvp`` so that ``defjvp`` can be used as a decorator.

    Examples:

      >>> @jax.custom_jvp
      ... def f(x, y):
      ...   return jnp.sin(x) * y
      ...
      >>> @f.defjvp
      ... def f_jvp(primals, tangents):
      ...   x, y = primals
      ...   x_dot, y_dot = tangents
      ...   primal_out = f(x, y)
      ...   tangent_out = jnp.cos(x) * x_dot * y + jnp.sin(x) * y_dot
      ...   return primal_out, tangent_out

      >>> x = jnp.float32(1.0)
      >>> y = jnp.float32(2.0)
      >>> with jnp.printoptions(precision=2):
      ...   print(jax.value_and_grad(f)(x, y))
      (Array(1.68, dtype=float32), Array(1.08, dtype=float32))
    """
    self.jvp = jvp
    self.symbolic_zeros = symbolic_zeros
    return jvp

  def defjvps(self, *jvps: Callable[..., ReturnValue] | None) -> None:
    """Convenience wrapper for defining JVPs for each argument separately.

    This convenience wrapper cannot be used together with ``nondiff_argnums``.

    Args:
      *jvps: a sequence of functions, one for each positional argument of the
        :class:`~jax.custom_jvp` function. Each function takes as arguments
        the tangent value for the corresponding primal input, the primal
        output, and the primal inputs. See the example below.

    Returns:
      None.

    Examples:

      >>> @jax.custom_jvp
      ... def f(x, y):
      ...   return jnp.sin(x) * y
      ...
      >>> f.defjvps(lambda x_dot, primal_out, x, y: jnp.cos(x) * x_dot * y,
      ...           lambda y_dot, primal_out, x, y: jnp.sin(x) * y_dot)

      >>> x = jnp.float32(1.0)
      >>> y = jnp.float32(2.0)
      >>> with jnp.printoptions(precision=2):
      ...   print(jax.value_and_grad(f)(x, y))
      (Array(1.68, dtype=float32), Array(1.08, dtype=float32))
    """
    if self.nondiff_argnums:
      raise TypeError("Can't use ``defjvps`` with ``nondiff_argnums``.")

    def jvp(primals, tangents):
      primal_out = self(*primals)
      zeros = _zeros_like_pytree(primal_out)
      all_tangents_out = [jvp(t, primal_out, *primals) if jvp else zeros
                          for t, jvp in zip(tangents, jvps)]
      tangent_out = tree_map(_sum_tangents, primal_out, *all_tangents_out)
      return primal_out, tangent_out

    self.defjvp(jvp)

  @partial(traceback_util.api_boundary,
           repro_api_name="jax.custom_jvp.__call__")
  def __call__(self, *args: Any, **kwargs: Any) -> ReturnValue:
    if config.custom_jvp3.value:
      from jax._src.hijax import custom_jvp3  # pyrefly: ignore[missing-import]
      f = custom_jvp3(self.fun, self.nondiff_argnums)
      f.defjvp(self.jvp, symbolic_zeros=self.symbolic_zeros)
      return f(*args, **kwargs)
    debug = debug_info("custom_jvp fun", self.fun, args, kwargs,
                       static_argnums=self.nondiff_argnums)
    primal_name = debug.func_name
    if not self.jvp:
      msg = f"No JVP defined for custom_jvp function {primal_name} using defjvp."
      raise AttributeError(msg)

    try:
      args = resolve_kwargs(self.fun, args, kwargs)
    except TypeError as e:
      raise TypeError(
          "The input arguments to the custom_jvp-decorated function "
          f"{primal_name} could not be resolved to positional-only arguments. "
          f"Binding failed with the error:\n{e}"
      ) from e

    if self.nondiff_argnums:
      args = tuple(_stop_gradient(x) if i in self.nondiff_argnums else x
                   for i, x in enumerate(args))
      diff_argnums = [i for i in range(len(args)) if i not in self.nondiff_argnums]
      f_, dyn_args = argnums_partial(lu.wrap_init(self.fun, debug_info=debug),
                                     diff_argnums, args,
                                     require_static_args_hashable=False)
      static_args = [args[i] for i in self.nondiff_argnums]
      diff_args = [args[i] for i, a in enumerate(args) if i not in self.nondiff_argnums]
      debug_jvp = debug_info("custom_jvp jvp", self.jvp,
                             (*static_args, diff_args, diff_args),
                             {},
                             static_argnums=tuple(range(len(static_args))))
      jvp = prepend_static_args(lu.wrap_init(self.jvp,
                                             debug_info=debug_jvp), static_args)
    else:
      f_, dyn_args = lu.wrap_init(self.fun, debug_info=debug), args
      debug_jvp = debug_info("custom_jvp jvp", self.jvp,
                             (args, args),
                             {})
      jvp = lu.wrap_init(self.jvp, debug_info=debug_jvp)
    args_flat, in_tree = tree_flatten(dyn_args)
    flat_fun, out_type1 = _flatten_fun_nokwargs(f_, in_tree)
    flat_jvp, out_type2 = _flatten_jvp(jvp, primal_name, debug_jvp.func_name,
                                       in_tree, out_type1)
    out_flat = custom_jvp_call_p.bind(*args_flat, subfuns=(flat_fun, flat_jvp),
                                      symbolic_zeros=self.symbolic_zeros)
    _, (out_tree, _, _) = lu.merge_linear_aux(out_type1, out_type2)
    return tree_unflatten(out_tree, out_flat)

@partial(lu.transformation_with_aux2, use_eq_store=True)
def _flatten_jvp(f, store, primal_name, jvp_name, in_tree, maybe_out_type, *args):
  primals_in, tangents_in = split_list(args, [len(args) // 2])
  py_primals = tree_unflatten(in_tree, primals_in)
  py_tangents = tree_unflatten(in_tree, tangents_in)
  pair_out = f(py_primals, py_tangents)
  if not isinstance(pair_out, (list, tuple)) or len(pair_out) != 2:
    msg = (f"Custom JVP rule {jvp_name} for function {primal_name} "
           "must produce a pair (list or tuple of length two) representing "
           f"primal and tangent outputs, but got {pair_out}.")
    raise TypeError(msg)
  py_primals_out, py_tangents_out = pair_out
  primals_out, out_tree = tree_flatten(py_primals_out)
  tangents_out, out_tree2 = tree_flatten(py_tangents_out)
  primal_avals = [core.typeof(x) for x in primals_out]
  if out_tree != out_tree2:
    msg = (f"Custom JVP rule {jvp_name} for function {primal_name} must "
           "produce primal and tangent outputs with equal container (pytree) "
           f"structures, but got {out_tree} and {out_tree2} respectively.")
    raise TypeError(msg)
  # If the primal function already ran, check out_tree agreement.
  try: out_type_ = maybe_out_type()
  except lu.StoreException: out_type_ = None
  if out_type_ is not None:
    out_tree_, primal_avals_, () = out_type_
    ty_tree  = tree_unflatten(out_tree , [a.str_short() for a in primal_avals])
    ty_tree_ = tree_unflatten(out_tree_, [a.str_short() for a in primal_avals_])
    if out_tree_ != out_tree:
      m = (f"Custom JVP rule {jvp_name} for function {primal_name} must "
           "produce a pair (list or tuple of length two) "
           "where the first element represents the primal output "
           "(equal in value to the output of the custom_jvp-decorated function "
           f"{primal_name}, "
           "and in particular of the same container/pytree structure), but "
           "instead the JVP rule output's first element had container/pytree "
           "structure:\n"
           f"""    {str(ty_tree ).replace("'", "")}\n"""
           f"while the custom_jvp-decorated function {primal_name} had output "
           "container/pytree structure:\n"
           f"""    {str(ty_tree_).replace("'", "")}.""")
      raise TypeError(m)
    if not all(map(core.typematch, primal_avals, primal_avals_)):
      m = (f"Custom JVP rule {jvp_name} for function {primal_name} must "
           "produce a pair (list or tuple of length two) "
           "where the first element represents the primal output "
           "(equal in value to the output of the custom_jvp-decorated function "
           f"{primal_name}, "
           "and in particular with leaves of the same shape/dtype), but "
           "instead the JVP rule output's first element had shapes/dtypes of:\n"
           f"""    {str(ty_tree ).replace("'", "")}\n"""
           f"while the custom_jvp-decorated function {primal_name} had output "
           "shapes/dtypes of:\n"
           f"""    {str(ty_tree_).replace("'", "")}""")
      raise TypeError(m)
  primal_avals_out = [core.typeof(x).strip_weak_type() for x in primals_out]
  expected_tangent_avals_out = [
    core.typeof(x).strip_weak_type().to_tangent_aval()
    for x in primals_out]
  tangent_avals_out = [core.typeof(t).strip_weak_type()
                       if type(t) is not SymbolicZero else t.aval.strip_weak_type()
                       for t in tangents_out]
  if not all(map(core.typematch, expected_tangent_avals_out, tangent_avals_out)):
    if len(expected_tangent_avals_out) == 1:
      (av_p,), (av_et,), (av_t,) = primal_avals_out, expected_tangent_avals_out, tangent_avals_out
      msg = ("Custom JVP rule must produce primal and tangent outputs with "
             "corresponding shapes and dtypes. Expected {} (tangent type of {}) but got {}.")
      raise TypeError(msg.format(av_et.str_short(), av_p.str_short(), av_t.str_short()))
    else:
      msg = ("Custom JVP rule must produce primal and tangent outputs with "
             "corresponding shapes and dtypes, but got:\n{}")
      disagreements = (
          f"  primal {av_p.str_short()} with tangent {av_t.str_short()}, expecting tangent {av_et}"
          for av_p, av_et, av_t in zip(primal_avals_out, expected_tangent_avals_out, tangent_avals_out)
          if av_et != av_t)
      raise TypeError(msg.format('\n'.join(disagreements)))
  store.store((out_tree, primal_avals, ()))
  return primals_out + tangents_out

class CustomJVPCallPrimitive(core.Primitive):
  multiple_results = True
  skip_canonicalization = True

  def bind_with_trace(self, trace, args, avals, params, /):
    params = dict(params)
    fun, jvp = params.pop('subfuns')
    return trace.process_custom_jvp_call(self, fun, jvp, args, **params)

  def impl(self, fun, _, *args):
    raise NotImplementedError

  def is_high(self, *_, call_jaxpr, **__):
    return call_jaxpr.is_high

  def to_lojax(self, *hi_args, call_jaxpr: core.Jaxpr, **params):
    return pe._lower_and_eval(pe.eval_jaxpr_p, call_jaxpr, hi_args)

  def get_bind_params(self, params):
    new_params = dict(params)
    call_jaxpr: core.Jaxpr = new_params.pop('call_jaxpr')
    num_consts: int = new_params.pop('num_consts')
    jvp_jaxpr_fun = new_params.pop('jvp_jaxpr_fun')
    fun = lu.wrap_init(core.jaxpr_as_fun(call_jaxpr),
                       debug_info=call_jaxpr.debug_info)
    jvp = lift_jvp(num_consts, jvp_jaxpr_fun)
    new_params['subfuns'] = (fun, jvp)
    return new_params

def lift_jvp(num_consts: int, jvp_jaxpr_fun: lu.WrappedFun) -> lu.WrappedFun:
  def jvp(*xs):
    n, ragged = divmod(len(xs), 2)
    assert not ragged
    primals, tangents = xs[num_consts:n], xs[n+num_consts:]
    zeros = [type(t) is SymbolicZero for t in tangents]
    jvp_jaxpr, jvp_consts, out_zeros = jvp_jaxpr_fun.call_wrapped(*zeros)
    nonzero_tangents = [t for t in tangents if type(t) is not SymbolicZero]
    out = core.eval_jaxpr(jvp_jaxpr, jvp_consts, *primals, *nonzero_tangents)
    out_primals, nz_out_tangents = split_list(out, [len(out_zeros)])
    nz_out_tangents_ = iter(nz_out_tangents)
    out_tangents = [SymbolicZero(core.typeof(p).to_tangent_aval())
                    if z else next(nz_out_tangents_)
                    for p, z in zip(out_primals, out_zeros)]
    assert next(nz_out_tangents_, None) is None
    return [*out_primals, *out_tangents]
  return lu.wrap_init(jvp, debug_info=jvp_jaxpr_fun.debug_info)

custom_jvp_call_p = CustomJVPCallPrimitive('custom_jvp_call')

def _custom_jvp_call_typecheck(_, *in_avals, call_jaxpr, jvp_jaxpr_fun,
                               num_consts, symbolic_zeros):
  # TODO(mattjj): could do more checking here...
  del in_avals, jvp_jaxpr_fun, num_consts
  disallowed_effects = effects.custom_derivatives_allowed_effects.filter_not_in(call_jaxpr.effects)
  if disallowed_effects:
    raise NotImplementedError(
        f'Effects not supported in `custom_jvp`: {disallowed_effects}')
  return call_jaxpr.out_avals, core.positional_effects(call_jaxpr)
core.custom_typechecks[custom_jvp_call_p] = _custom_jvp_call_typecheck

def _custom_jvp_vjp_call_lowering(ctx: mlir.LoweringRuleContext, *args,
                                  call_jaxpr: core.Jaxpr, **_):
  consts = mlir.ir_consts(
      call_jaxpr.consts, [v.aval for v in call_jaxpr.constvars])
  out, tokens = mlir.jaxpr_subcomp(ctx.module_context, call_jaxpr,
                                   ctx.name_stack, ctx.tokens_in, consts,
                                   *args, dim_var_values=ctx.dim_var_values,
                                   const_lowering=ctx.const_lowering,
                                   outer_traceback=ctx.traceback)
  ctx.set_tokens_out(tokens)
  return out
mlir.register_lowering(custom_jvp_call_p, _custom_jvp_vjp_call_lowering)

def _custom_jvp_call_transpose_fancy(cts, *args, call_jaxpr, **_):
  return ad.backward_pass3(call_jaxpr, False, call_jaxpr.consts, args, cts)
ad.fancy_transposes[custom_jvp_call_p] = _custom_jvp_call_transpose_fancy

@weakref_lru_cache
def _cached_closed_call_dce_instantiate(jaxpr_: core.Jaxpr,
                                        used_outputs: tuple[bool, ...]
                                        ) -> tuple[core.Jaxpr, list[bool]]:
  # dce_jaxpr and replace preserve attached consts.
  return pe.dce_jaxpr(
      jaxpr_.replace(debug_info=jaxpr_.debug_info.with_unknown_names()),
      used_outputs, True)

def _custom_jvp_call_dce(
    used_outs: Sequence[bool], live_ins: Sequence[bool], eqn: core.JaxprEqn
) -> tuple[list[bool], core.JaxprEqn | None]:
  if not any(used_outs) and not pe.has_effects(eqn, live_ins):
    return [False] * len(eqn.invars), None

  call_jaxpr = eqn.params["call_jaxpr"]
  jvp_jaxpr_fun = eqn.params["jvp_jaxpr_fun"]
  # We must set instantiate=True because some inputs that are unused by the
  # DCE'ed primal might be used in the JVP rule.
  dce_call_jaxpr, used_ins = _cached_closed_call_dce_instantiate(
      call_jaxpr, tuple(used_outs))
  assert all(used_ins)

  @pe._memoize
  def dce_jvp_jaxpr_thunk(*in_zeros):
    jvp_jaxpr, consts, out_zeros = jvp_jaxpr_fun.call_wrapped(*in_zeros)
    sz = eqn.params["symbolic_zeros"]
    nz_used_outs = [u for u, z in zip(used_outs, out_zeros) if not z] if sz else used_outs
    dce_jvp_jaxpr, _ = pe.dce_jaxpr(jvp_jaxpr, [*used_outs, *nz_used_outs], True)
    dce_out_zeros = [v for used, v in zip(used_outs, out_zeros) if used]
    return dce_jvp_jaxpr, consts, dce_out_zeros

  outvars = [v for used, v in zip(used_outs, eqn.outvars) if used]
  new_params = dict(
      eqn.params,
      call_jaxpr=dce_call_jaxpr,
      jvp_jaxpr_fun=lu.wrap_init(dce_jvp_jaxpr_thunk,
                                 debug_info=jvp_jaxpr_fun.debug_info)
  )
  new_eqn = pe.new_jaxpr_eqn(
      eqn.invars, outvars, eqn.primitive, new_params,
      core.eqn_effects(dce_call_jaxpr, eqn.invars),
      eqn.source_info, eqn.ctx)
  return used_ins, new_eqn
pe.dce_rules[custom_jvp_call_p] = _custom_jvp_call_dce


def _custom_jvp_call_pp_rule(eqn: core.JaxprEqn,
                             context: core.JaxprPpContext,
                             settings: core.JaxprPpSettings) -> core.pp.Doc:
  params = dict(eqn.params)
  if not params["num_consts"]:
    params.pop("num_consts")
  params["jvp"] = params.pop("jvp_jaxpr_fun").debug_info.func_name
  names = sorted(params)
  params["name"] = params["call_jaxpr"].debug_info.func_name
  return core._pp_eqn(eqn.replace(params=params), context, settings,
                      params=["name"] + names)


core.pp_eqn_rules[custom_jvp_call_p] = _custom_jvp_call_pp_rule

### VJPs

@custom_api_util.register_custom_decorator_type
class custom_vjp[ReturnValue]:
  """Set up a JAX-transformable function for a custom VJP rule definition.

  This class is meant to be used as a function decorator. Instances are
  callables that behave similarly to the underlying function to which the
  decorator was applied, except when a reverse-mode differentiation
  transformation (like :py:func:`jax.grad`) is applied, in which case a custom
  user-supplied VJP rule function is used instead of tracing into and performing
  automatic differentiation of the underlying function's implementation. There
  is a single instance method, :py:func:`~jax.custom_vjp.defvjp`, which may be
  used to define the custom VJP rule.

  This decorator precludes the use of forward-mode automatic differentiation.

  For example::

    @jax.custom_vjp
    def f(x, y):
      return jnp.sin(x) * y

    def f_fwd(x, y):
      return f(x, y), (jnp.cos(x), jnp.sin(x), y)

    def f_bwd(res, g):
      cos_x, sin_x, y = res
      return (cos_x * g * y, sin_x * g)

    f.defvjp(f_fwd, f_bwd)

  For a more detailed introduction, see the tutorial_.

  .. _tutorial: https://docs.jax.dev/en/latest/301/custom-jvp-vjp.html
  """

  def __init__(self,
               fun: Callable[..., ReturnValue],
               nondiff_argnums: Sequence[int] = (),
               nondiff_argnames: Sequence[str] = ()):
    update_wrapper(self, fun)
    self.fun = fun

    nondiff_argnums_: set[int] = set()
    if nondiff_argnames:
      sig = fun_signature(self.fun)
      assert sig is not None
      inferred_nondiff_argnums, _ = infer_argnums_and_argnames(
          sig, None, nondiff_argnames
      )
      nondiff_argnums_.update(inferred_nondiff_argnums)

    if nondiff_argnums:
      nondiff_argnums_.update(nondiff_argnums)

    self.nondiff_argnums = tuple(sorted(nondiff_argnums_))
    self.fwd: Callable[..., tuple[ReturnValue, Any]] | None = None
    self.bwd: Callable[..., tuple[Any, ...]] | None = None
    self.symbolic_zeros = False
    self.optimize_remat = False
    self.with_logs = False
    self.with_accums = False
    self.generic_batching = False
    self.remat_rules: tuple[Callable, Callable, Callable, bool] | None = None

  __getattr__ = custom_api_util.forward_attr

  def defvjp(self,
             fwd: Callable[..., tuple[ReturnValue, Any]],
             bwd: Callable[..., tuple[Any, ...]],
             symbolic_zeros: bool = False,
             optimize_remat: bool = False,
             generic_batching: bool = False,
             ) -> None:
    """Define a custom VJP rule for the function represented by this instance.

    Args:
      fwd: a Python callable representing the forward pass of the custom VJP
        rule. When there are no ``nondiff_argnums``, the ``fwd`` function has
        the same input signature as the underlying primal function. It should
        return as output a pair, where the first element represents the primal
        output and the second element represents any "residual" values to store
        from the forward pass for use on the backward pass by the function
        ``bwd``. Input arguments and elements of the output pair may be arrays
        or nested tuples/lists/dicts thereof.
      bwd: a Python callable representing the backward pass of the custom VJP
        rule. When there are no ``nondiff_argnums``, the ``bwd`` function takes
        two arguments, where the first is the "residual" values produced on the
        forward pass by ``fwd``, and the second is the output cotangent with the
        same structure as the primal function output. The output of ``bwd`` must
        be a tuple of length equal to the number of arguments of the primal
        function, and the tuple elements may be arrays or nested
        tuples/lists/dicts thereof so as to match the structure of the primal
        input arguments.
      symbolic_zeros: boolean, determining whether to indicate symbolic zeros
        to the ``fwd`` and ``bwd`` rules. Enabling this option allows custom
        derivative rules to detect when certain inputs, and when certain
        output cotangents, are not involved in differentiation. If ``True``:

        * ``fwd`` must accept, in place of each leaf value ``x`` in
          the pytree comprising an argument to the original function,
          an object (of type
          ``jax.custom_derivatives.CustomVJPPrimal``) with two
          attributes instead: ``value`` and ``perturbed``. The
          ``value`` field is the original primal argument, and
          ``perturbed`` is a boolean.  The ``perturbed`` bit indicates
          whether the argument is involved in differentiation (i.e.,
          if it is ``False``, then the corresponding Jacobian "column"
          is zero).

        * ``bwd`` will be passed objects representing static symbolic zeros in
          its cotangent argument in correspondence with unperturbed values;
          otherwise, only standard JAX types (e.g. array-likes) are passed.

        Setting this option to ``True`` allows these rules to detect whether
        certain inputs and outputs are not involved in differentiation, but at
        the cost of special handling. For instance:

        * The signature of ``fwd`` changes, and the objects it is passed cannot
          be output from the rule directly.

        * The ``bwd`` rule is passed objects that are not entirely array-like,
          and that cannot be passed to most ``jax.numpy`` functions.

        * Any custom pytree nodes involved in the primal function's arguments
          must accept, in their unflattening functions, the two-field record
          objects that are given as input leaves to the ``fwd`` rule.

        Default ``False``.
      optimize_remat: boolean, an experimental flag to enable an automatic
        optimization when this function is used under :func:`jax.remat`. This
        will be most useful when the ``fwd`` rule is an opaque call such as a
        Pallas kernel or a custom call. Default ``False``.
      generic_batching: boolean, determining how ``vmap`` handles a primal
        output that is shared across the batch. By default, when such an
        output has a nonzero cotangent and ``bwd`` produces a batched
        cotangent for an unbatched primal input, that raises an error, since
        summing over the batch would count the shared output's cotangent once
        per example. If ``True``, ``bwd`` instead runs with the shared output's
        whole cotangent given to the first example, which is right for a
        correct rule but does ``bwd``'s work on that cotangent once per
        example. Likewise, an output of ``fwd`` that is batched where the
        primal output is not takes the first example's value. See https://docs.jax.dev/en/latest/301/custom-derivatives.html#jax-301-vmap-shared-outputs.
        Default ``False``.

    Returns:
      None.

    Examples:

      >>> @jax.custom_vjp
      ... def f(x, y):
      ...   return jnp.sin(x) * y
      ...
      >>> def f_fwd(x, y):
      ...   return f(x, y), (jnp.cos(x), jnp.sin(x), y)
      ...
      >>> def f_bwd(res, g):
      ...   cos_x, sin_x, y = res
      ...   return (cos_x * g * y, sin_x * g)
      ...
      >>> f.defvjp(f_fwd, f_bwd)

      >>> x = jnp.float32(1.0)
      >>> y = jnp.float32(2.0)
      >>> with jnp.printoptions(precision=2):
      ...   print(jax.value_and_grad(f)(x, y))
      (Array(1.68, dtype=float32), Array(1.08, dtype=float32))
    """
    self.fwd = fwd
    self.bwd = bwd
    self.symbolic_zeros = symbolic_zeros
    self.optimize_remat = optimize_remat
    self.with_accums = False
    self.generic_batching = generic_batching
    if self.symbolic_zeros and self.optimize_remat:
      raise NotImplementedError(
          "remat optimization for custom_vjp does not support symbolic zeros")

  def defvjp_with_logs(self,
                       fwd: Callable[..., tuple[ReturnValue, Any]],
                       bwd: Callable[..., tuple[tuple[Any, ...], dict | None]],
                       symbolic_zeros: bool = False,
                       optimize_remat: bool = False,
                       generic_batching: bool = False,
                       ) -> None:
    """Like :py:func:`~jax.custom_vjp.defvjp`, but ``bwd`` can also log.

    The only difference from ``defvjp`` is the return convention of ``bwd``:
    it must return a pair ``(in_cts, logs)``, where ``in_cts`` is the usual
    tuple of cotangents (one entry per primal argument), and ``logs`` is a
    dict of named pytrees to log out of the backward pass, or ``None`` to log
    nothing. To receive the logs, apply the VJP function via its
    ``with_logs`` method: ``f_vjp.with_logs(out_ct)`` returns a pair
    ``(arg_cts, logs)``. Logging is drop-by-default: a plain ``f_vjp(out_ct)``
    call ignores the logs, and under ``jit`` the logging computation is
    dead-code-eliminated. Log keys must be unique within a backward pass;
    repeated keys raise a ``ValueError``, even when the logs are ignored.
    """
    self.defvjp(fwd, bwd, symbolic_zeros=symbolic_zeros,
                optimize_remat=optimize_remat, generic_batching=generic_batching)
    self.with_logs = True

  def defvjp_with_accums(self,
                         fwd: Callable[..., tuple[ReturnValue, Any]],
                         bwd: Callable[..., Any],  # returns None or a dict of logs
                         symbolic_zeros: bool = False,
                         optimize_remat: bool = False,
                         generic_batching: bool = False,
                         ) -> None:
    """Like :py:func:`~jax.custom_vjp.defvjp`, but ``bwd`` takes gradient
    accumulators.

    Rather than returning a tuple of cotangents, ``bwd`` takes one gradient
    accumulator per argument of the primal function (matching its pytree
    structure) after its two usual arguments, as in ``bwd(res, g, *arg_accums)``,
    and pushes each cotangent contribution into the corresponding accumulator.
    It returns ``None``, or a dict of backward-pass logs as with
    :py:func:`~jax.custom_vjp.defvjp_with_logs`. The accumulators are the
    ``ValAccum``, ``RefAccum`` and ``NullAccum`` objects that the ``vjp_bwd``
    method of a hijax primitive receives (see :ref:`jax-301-vjp-bwd-accums`).
    Under transformations that wrap the backward rule, such as :func:`jax.vmap`
    applied before differentiation, ``bwd`` receives a ``ValAccum`` in place of
    a ``RefAccum``, and the result is added to the ``Ref`` afterward.
    """
    self.defvjp(fwd, bwd, symbolic_zeros=symbolic_zeros,
                optimize_remat=optimize_remat, generic_batching=generic_batching)
    self.with_accums = True

  def defremat(self,
               fwd: Callable[..., tuple[ReturnValue, Any]],
               rem: Callable[..., tuple[ReturnValue, Any]],
               bwd: Callable[..., tuple[Any, ...]],
               generic_batching: bool = False,
               ) -> None:
    """Define custom rematerialization rules for this function.

    The rules apply when this function is differentiated in reverse mode
    under :func:`jax.remat`, where by default the ``fwd`` rule given to
    :py:func:`~jax.custom_vjp.defvjp` is rematerialized according to the
    ambient checkpoint policy. Outside of :func:`jax.remat`, unless a VJP is
    also defined with ``defvjp``, the rules define the VJP too, by running
    ``rem`` right after ``fwd`` on the forward pass. Requires the
    ``jax_custom_vjp3`` and ``jax_remat3`` implementations.

    Args:
      fwd: forward-pass rule, with the same signature as the ``fwd`` rule of
        :py:func:`~jax.custom_vjp.defvjp`: it takes the arguments of the
        primal function, and returns the primal output paired with residuals
        to save for the backward pass.
      rem: rematerialization rule, of signature ``rem(res, *args) -> (out,
        res2)``. On the backward pass it receives the residuals saved by
        ``fwd`` and the arguments of the primal function, and recomputes the
        primal output paired with the residuals ``res2`` for ``bwd``.
      bwd: backward-pass rule, as in :py:func:`~jax.custom_vjp.defvjp`,
        receiving ``res2`` as its residuals.
      generic_batching: as in :py:func:`~jax.custom_vjp.defvjp`.

    Returns:
      None.
    """
    self.remat_rules = (fwd, rem, bwd, False)
    self.generic_batching = generic_batching

  def defremat_with_logs(self,
                         fwd: Callable[..., tuple[ReturnValue, Any]],
                         rem: Callable[..., tuple[ReturnValue, Any]],
                         bwd: Callable[..., tuple[tuple[Any, ...], dict | None]],
                         generic_batching: bool = False,
                         ) -> None:
    """Like :py:func:`~jax.custom_vjp.defremat`, but ``bwd`` can also log.

    As with :py:func:`~jax.custom_vjp.defvjp_with_logs`, ``bwd`` must return a
    pair ``(in_cts, logs)``.
    """
    self.remat_rules = (fwd, rem, bwd, True)
    self.generic_batching = generic_batching

  @partial(traceback_util.api_boundary,
           repro_api_name="jax.custom_vjp.__call__")
  def __call__(self, *args: Any, **kwargs: Any) -> ReturnValue:
    if config.custom_vjp3.value or self.generic_batching:
      from jax._src.hijax import custom_vjp3  # pyrefly: ignore[missing-import]
      f = custom_vjp3(self.fun, self.nondiff_argnums)
      if self.fwd is not None:
        (f.defvjp_with_accums if self.with_accums else
         f.defvjp_with_logs if self.with_logs else f.defvjp)(
            self.fwd, self.bwd, symbolic_zeros=self.symbolic_zeros,
            optimize_remat=self.optimize_remat)
      f.remat_rules = self.remat_rules
      f.generic_batching = self.generic_batching
      return f(*args, **kwargs)
    if self.remat_rules is not None and not (self.fwd and self.bwd):
      raise NotImplementedError(
          "custom_vjp.defremat requires the jax_custom_vjp3 implementation, "
          "enabled with jax.config.update('jax_custom_vjp3', True).")
    debug_fun = debug_info("custom_vjp fun", self.fun, args, kwargs,
                           static_argnums=self.nondiff_argnums)
    if not self.fwd or not self.bwd:
      msg = f"No VJP defined for custom_vjp function {debug_fun.func_name} using defvjp."
      raise AttributeError(msg)

    try:
      args = resolve_kwargs(self.fun, args, kwargs)
    except TypeError as e:
      raise TypeError(
          "The input arguments to the custom_vjp-decorated function "
          f"{debug_fun.func_name} could not be resolved to positional-only "
          f"arguments. Binding failed with the error:\n{e}"
      ) from e

    debug_fwd = debug_info("custom_vjp fwd", self.fwd, args, kwargs,
                           static_argnums=self.nondiff_argnums)
    # TODO(necula): figure out how to construct the debug_bwd args
    debug_bwd = debug_info("custom_vjp bwd", self.bwd, args, {})
    if self.optimize_remat:
      fwd = optimize_remat_of_custom_vjp_fwd(
          self.fun, debug_fun, self.fwd, debug_fwd,
          nondiff_argnums=self.nondiff_argnums,
          symbolic_zeros=self.symbolic_zeros)
    else:
      fwd = self.fwd
    if self.nondiff_argnums:
      for i in self.nondiff_argnums: _check_for_tracers(args[i])
      dyn_argnums = [i for i in range(len(args)) if i not in self.nondiff_argnums]
      f_, dyn_args = argnums_partial(
          lu.wrap_init(self.fun, debug_info=debug_fun), dyn_argnums,
          args, require_static_args_hashable=False)
      static_args = [args[i] for i in self.nondiff_argnums]
      fwd_, _ = argnums_partial(lu.wrap_init(fwd, debug_info=debug_fwd),
                                dyn_argnums, args,
                                require_static_args_hashable=False)
      bwd = prepend_static_args(lu.wrap_init(self.bwd, debug_info=debug_bwd),
                                static_args)
    else:
      f_, dyn_args = lu.wrap_init(self.fun, debug_info=debug_fun), args
      fwd_ = lu.wrap_init(fwd, debug_info=debug_fwd)
      bwd = lu.wrap_init(self.bwd, debug_info=debug_bwd)
    args_flat, in_tree = tree_flatten(dyn_args)
    in_avals = tuple(core.typeof(x) for x in args_flat)
    if config.mutable_array_checks.value:
      _check_for_aliased_refs(self.fun, (), debug_fun, dyn_args)
      f_ = _check_primal_refs(f_)
    flat_fun, out_type = _flatten_fun_nokwargs(f_, in_tree)
    flat_fwd, out_trees = _flatten_fwd(
        fwd_, self.symbolic_zeros, debug_fun, debug_fwd, in_tree, out_type)
    if self.with_accums:
      bwd_ = self.bwd
      if self.nondiff_argnums:
        bwd_ = partial(self.bwd, *[args[i] for i in self.nondiff_argnums])
      flat_bwd = lu.wrap_init(
          _FlatBwdWithAccums(bwd_, in_tree, in_avals, out_trees,
                             self.symbolic_zeros),
          debug_info=debug_bwd)
    else:
      flat_bwd = _flatten_bwd(bwd, in_tree, in_avals, out_trees, self.fun,
                              self.symbolic_zeros, self.with_logs)
    out_flat = custom_vjp_call_p.bind(*args_flat, subfuns=(flat_fun, flat_fwd, flat_bwd),
                                      out_trees=out_trees,
                                      symbolic_zeros=self.symbolic_zeros)
    _, (out_tree, _, _) = lu.merge_linear_aux(out_type, out_trees)
    return tree_unflatten(out_tree, out_flat)

@lu.transformation2
def _check_primal_refs(f: Callable, *args):
  out = f(*args)
  _check_for_returned_refs(f, out, 'primal', [], 0)
  return out

def _check_for_aliased_refs(
    f: Callable, nondiff_argnums: Sequence[int], debug: core.DebugInfo, args):
  argnums = [x for i, arg in enumerate(args)
             for x in [i] * tree_structure(arg).num_leaves]
  leaves = tree_leaves(args)
  refs: dict[int, int] = {}
  for i, (argnum, x) in enumerate(zip(argnums, leaves)):
    if argnum in nondiff_argnums: continue
    x = x.value if isinstance(x, CustomVJPPrimal) else x
    if (isinstance((a := core.typeof(x)), AbstractRef) and
        (dup_idx := refs.setdefault(id(core.get_referent(x)), i)) != i):
      arg_names = debug.safe_arg_names(len(leaves))
      raise ValueError(
          "only one reference to a mutable array may be passed as an argument "
          f"to a function, but custom_vjp function {f} got the same mutable "
          f"array reference of type {a.str_short()} at {arg_names[dup_idx]} and"
          f" {arg_names[i]}.")

def _check_for_returned_refs(f, out, kind, args, after_idx):
  args = [x.value if isinstance(x, CustomVJPPrimal) else x for x in args]
  ids = {id(x) for x in args if isinstance(core.typeof(x), AbstractRef)}
  leaves = tree_leaves_with_path(out)
  for i, (path, leaf) in enumerate(leaves):
    if isinstance((a := core.typeof(leaf)), AbstractRef):
      loc = f' at output tree path {keystr(path)}' if path else ''
      if i < after_idx:
        raise ValueError(f"custom_vjp {kind} function {f} returned a mutable "
                         f"array reference of type {a.str_short()}{loc}, "
                         "but mutable array references cannot be returned there.")
      if id(leaf) not in ids:
        raise ValueError(f"custom_vjp {kind} function {f} returned a mutable "
                         f"array reference of type {a.str_short()}{loc} "
                         "that was not an argument.")

@dataclasses.dataclass(slots=True)
class CustomVJPPrimal:
  """Primal to a ``custom_vjp``'s forward rule when ``symbolic_zeros`` is set"""
  value: Any
  perturbed: bool

def custom_vjp_primal_tree_values(tree):
  """Strips away perturbation information from forward rule arguments.

  This is a helper function for user with the ``symbolic_zeros`` option to
  the ``defvjp`` method of a ``custom_vjp``-decorated function.

  In ``symbolic_zeros`` mode, the custom forward rule receives arguments
  whose pytree leaves are records with a ``value`` attribute that carries
  the primal argument. This is a way to convert such argument trees back to
  their original form, replacing each such record with its carried value at
  each leaf.
  """
  def value(leaf):
    if type(leaf) is not CustomVJPPrimal:
      raise TypeError(f"unexpected leaf type {type(leaf)}")
    return leaf.value
  return tree_map(value, tree)

def _check_for_tracers(x):
  for leaf in tree_leaves(x):
    if isinstance(leaf, core.Tracer):
      msg = ("Found a JAX Tracer object passed as an argument to a custom_vjp "
            "function in a position indicated by nondiff_argnums as "
            "non-differentiable. Tracers cannot be passed as non-differentiable "
            "arguments to custom_vjp functions; instead, nondiff_argnums should "
            "only be used for arguments that can't be or contain JAX tracers, "
            "e.g. function-valued arguments. In particular, array-valued "
            "arguments should typically not be indicated as nondiff_argnums.")
      raise UnexpectedTracerError(msg)

@partial(lu.transformation_with_aux2, use_eq_store=True)
def _flatten_fwd(f: Callable, store: lu.EqualStore,
                 symbolic_zeros: bool,
                 debug_primal: core.DebugInfo,
                 debug_fwd: core.DebugInfo,
                 in_tree: PyTreeDef, maybe_out_type, *args):
  primal_name = debug_primal.func_name if debug_primal else str(f)
  fwd_name = debug_fwd.func_name if debug_fwd else "<unknown>"
  if symbolic_zeros:
    args = tuple(CustomVJPPrimal(x, z) for x, z in zip(args[::2], args[1::2]))
  else:
    args = args[::2]
  py_args = tree_unflatten(in_tree, args)
  pair_out = f(*py_args)
  if not isinstance(pair_out, (list, tuple)) or len(pair_out) != 2:
    msg = (f"Custom VJP fwd rule {fwd_name} for function {primal_name} "
           "must produce a pair (list or tuple of length two) where the first "
           "element represents the primal output (equal to those of the "
           f"custom_vjp-decorated function {primal_name}) and the "
           "second element represents residuals (i.e. values stored from the "
           "forward pass for use on the backward pass), but "
           f"instead of a pair the fwd rule {fwd_name} produced {pair_out}.")
    raise TypeError(msg)
  py_primals_out, res = pair_out
  primals_out, out_tree = tree_flatten(py_primals_out)
  res, res_tree = tree_flatten(res)
  if config.mutable_array_checks.value:
    _check_for_returned_refs(f, pair_out, "fwd", args, out_tree.num_leaves)
  primal_avals = [core.typeof(x) for x in primals_out]
  # If the primal function already ran, check out_tree agreement.
  try: out_type_ = maybe_out_type()
  except lu.StoreException: out_type_ = None
  if out_type_ is not None:
    out_tree_, primal_avals_, () = out_type_
    ty_tree  = tree_unflatten(out_tree , [a.str_short() for a in primal_avals])
    ty_tree_ = tree_unflatten(out_tree_, [a.str_short() for a in primal_avals_])
    if out_tree_ != out_tree:
      m = (f"Custom VJP fwd rule {fwd_name} for function {primal_name} "
           "must produce a pair (list or tuple of length two) where the first "
           "element represents the primal output "
           "(equal to the output of the custom_vjp-decorated function "
           f"{primal_name}) and the "
           "second element represents residuals (i.e. values stored from the "
           "forward pass for use on the backward pass), but "
           "instead the fwd rule output's first element had container/pytree "
           "structure:\n"
           f"""    {str(ty_tree ).replace("'", "")}\n"""
           f"while the custom_vjp-decorated function {primal_name} had output "
           "container/pytree structure:\n"
           f"""    {str(ty_tree_).replace("'", "")}.""")
      raise TypeError(m)
    if not all(map(core.typematch, primal_avals, primal_avals_)):
      m = (f"Custom VJP fwd rule {fwd_name} for function {primal_name} must "
           "produce a pair (list or tuple of length two) "
           "where the first element represents the primal output "
           "(equal to the output of the custom_vjp-decorated function "
           f"{primal_name}) and the second element represents residuals "
           "(i.e. values stored from the forward pass for use on the "
           "backward pass), but "
           "instead the fwd rule output's first element had shapes/dtypes of:\n"
           f"""    {str(ty_tree ).replace("'", "")}\n"""
           f"while the custom_vjp-decorated function {primal_name} had output "
           "shapes/dtypes of:\n"
           f"""    {str(ty_tree_).replace("'", "")}""")
      raise TypeError(m)
  pruned_res, input_forwards = _filter_forwarded_inputs(res, args)  # prune
  store.store((out_tree, res_tree, input_forwards))
  return (*pruned_res, *primals_out)

def _filter_forwarded_inputs(outs, ins):
  idxs: dict[int, int] = {id(x): i for i, x in enumerate(ins)}
  return [o for o in outs if id(o) not in idxs], [idxs.get(id(o)) for o in outs]

class _FlatBwdWithAccums(ad.CustomBwdWithAccums):
  def __init__(self, bwd, in_tree, in_avals, out_trees, symbolic_zeros):
    self.bwd = bwd
    self.in_tree = in_tree
    self.ct_avals = [a.to_ct_aval() for a in in_avals]
    self.out_trees = out_trees
    self.symbolic_zeros = symbolic_zeros

  def call_with_accums(self, res, cts_out, accums):
    out_tree, res_tree, _ = self.out_trees()
    if not self.symbolic_zeros:
      cts_out = [zeros_like_aval(x.aval) if type(x) is SymbolicZero else x
                 for x in cts_out]
    accums = [acc if isinstance(acc, ad.GradAccum) else ad.NullAccum(a)
              for a, acc in zip(self.ct_avals, accums)]
    logs = self.bwd(tree_unflatten(res_tree, res),
                    tree_unflatten(out_tree, cts_out),
                    *tree_unflatten(self.in_tree, accums))
    if logs is not None and type(logs) is not dict:
      raise TypeError(
          "Custom VJP bwd rule was registered with defvjp_with_accums, and so "
          "must return None or a dict of backward-pass log entries, but got "
          f"{type(logs).__name__}.")
    return logs

  def __call__(self, *res_and_cts):
    _, res_tree, _ = self.out_trees()
    res, cts_out = split_list(res_and_cts, [res_tree.num_leaves])
    accums = [ad.NullAccum(a) if getattr(a, 'dtype', None) == dtypes.float0
              else ad.ValAccum(a) for a in self.ct_avals]
    logs = self.call_with_accums(res, cts_out, accums)
    cts_in = [acc.freeze() if isinstance(acc, ad.ValAccum) else Zero(acc.aval)
              for acc in accums]
    return cts_in, logs

@lu.transformation2
def _flatten_bwd(f: Callable,
                 in_tree: PyTreeDef,
                 in_avals: Sequence[core.AbstractValue],  # primal avals
                 out_trees: Callable[[], tuple[PyTreeDef, PyTreeDef, list[int | None]]],
                 primal_fun, symbolic_zeros: bool, with_logs: bool, *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])
  if not symbolic_zeros:
    cts_out = [zeros_like_aval(x.aval) if type(x) is SymbolicZero else x
               for x in cts_out]
  py_res = tree_unflatten(res_tree, res)
  py_cts_out = tree_unflatten(out_tree, cts_out)
  py_cts_in = f(py_res, py_cts_out)
  if with_logs:
    if not (isinstance(py_cts_in, (list, tuple)) and len(py_cts_in) == 2):
      raise TypeError(
          "Custom VJP bwd rule was registered with defvjp_with_logs and so "
          f"must produce a pair (in_cts, logs), but got {py_cts_in}.")
    py_cts_in, logs = py_cts_in
    if logs is not None and type(logs) is not dict:
      raise TypeError(
          "Custom VJP bwd rule was registered with defvjp_with_logs, and so "
          "the second element of the pair it returns must be None or a dict "
          f"of backward-pass log entries, but got {type(logs).__name__}.")
  else:
    logs = None
  if isinstance(py_cts_in, list) and len(py_cts_in) == len(treedef_children(in_tree)):
    py_cts_in = tuple(py_cts_in)
  # For each None in py_cts_in, indicating an argument for which the rule
  # produces no cotangent, we replace it with a pytree with the structure of the
  # corresponding subtree of in_tree and with leaves of a non-pytree sentinel
  # object, to be replaced with Nones in the final returned result.
  zero = object()  # non-pytree sentinel to replace Nones in py_cts_in
  dummy = tree_unflatten(in_tree, [object()] * in_tree.num_leaves)
  keypaths, _ = unzip2(tree_flatten_with_path(dummy)[0])
  cts_in_flat = []
  def append(x, d):
    num_leaves = len(tree_flatten(d)[0])
    if x is None and d is not None:
      cts_in_flat.extend([zero] * num_leaves)
    elif x is not None:
      cts_in_flat.extend([x] * num_leaves)
    return x
  try:
    if not isinstance(py_cts_in, tuple):
      raise ValueError
    tree_map(append, py_cts_in, dummy, is_leaf=lambda x: x is None)
  except ValueError:
    _, in_tree2 = tree_flatten(py_cts_in)
    msg = ("Custom VJP bwd rule must produce an output with the same container "
           "(pytree) structure as the args tuple of the primal function, "
           "and in particular must produce a tuple of length equal to the "
           "number of arguments to the primal function, but got bwd output "
           "structure {} for primal input structure {}.")
    raise TypeError(msg.format(in_tree2, in_tree)) from None
  results: list[Any] = []
  for kp, a, ct in zip(keypaths, in_avals, cts_in_flat):
    if ct is zero or getattr(a.to_ct_aval(), 'dtype') == dtypes.float0:
      results.append(Zero(a.to_ct_aval()))
    elif type(ct) is SymbolicZero:
      if not core.typecompat(a.to_ct_aval(), a_ := ct.aval):
        msg = ("Custom VJP bwd rule produced a SymbolicZero with a shape/dtype "
               "that does not match the corresponding input tangent shape/dtype: "
               f"at output{keystr(kp)} the SymbolicZero had shape/dtype "
               f"{a_.str_short()} while the "
               f"corresponding input had shape/dtype {a.str_short()}. "
               "Consider just returning a None here instead of a SymbolicZero "
               "object.")
        raise ValueError(msg)
      results.append(Zero(ct.aval))
    else:
      if (not config.disable_bwd_checks.value and
          not core.typecompat(a.to_ct_aval(), a_ := core.typeof(ct))
          and not _ref_typecompat(a.to_ct_aval(), a_)
          and not _temporary_dtype_exception(a.to_ct_aval(), a_)):
        primal_info = fun_sourceinfo(primal_fun)
        msg = (f"Custom VJP bwd rule attached to {primal_info} must produce an "
               "output with the same "
               "type as the args tuple of the primal function, but at "
               f"output{keystr(kp)} the bwd rule produced an output of "
               f"type {a_.str_short()} corresponding "
               f"to an input of type {a.str_short()}"
               f"{core.aval_mismatch_extra(a, a_)}")
        raise ValueError(msg)
      results.append(ct)
  return results, logs

def _ref_typecompat(a, a_):
  return (isinstance(a, AbstractRef) and
          core.typecompat(a.to_ct_aval().inner_aval, a_))

# TODO(mattjj): remove both these exceptions to cotangent compatibility check
def _temporary_dtype_exception(a, a_) -> bool:
  if isinstance(a, core.ShapedArray) and isinstance(a_, core.ShapedArray):
    return (a.shape == a_.shape and
            core.typematch(a, a_, no_dtype_check=True) and
            (dtypes.issubdtype(a_.dtype, dtypes.extended) or
             dtypes.issubdtype(a.dtype, dtypes.np.inexact)))
  return False


class CustomVJPCallPrimitive(core.Primitive):
  multiple_results = True
  skip_canonicalization = True

  def bind_with_trace(self, trace, args, avals, params, /):
    params = dict(params)
    fun, fwd, bwd = params.pop('subfuns')
    return trace.process_custom_vjp_call(self, fun, fwd, bwd, args, **params)

  def impl(self, fun, fwd, bwd, *args):
    raise NotImplementedError

  def is_high(self, *_, call_jaxpr, **__):
    return call_jaxpr.is_high

  def to_lojax(self, *hi_args, call_jaxpr: core.Jaxpr, **params):
    return pe._lower_and_eval(pe.eval_jaxpr_p, call_jaxpr, hi_args)

  def get_bind_params(self, params):
    new_params = dict(params)
    call_jaxpr: core.Jaxpr = 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.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)
    new_params['subfuns'] = (fun, fwd, bwd)
    return new_params

def lift_fwd(num_consts: int, fwd_jaxpr_thunk: lu.WrappedFun) -> lu.WrappedFun:
  def fwd(*args):
    vals, nonzeros = args[::2], args[1::2]
    assert len(vals) == len(nonzeros)
    _, primals = split_list(vals, [num_consts])
    const_nonzeros, in_nonzeros = split_list(nonzeros, [num_consts])
    if any(const_nonzeros): raise ad.CustomVJPException()
    fwd_jaxpr, fwd_consts = fwd_jaxpr_thunk.call_wrapped(*in_nonzeros)
    return core.eval_jaxpr(fwd_jaxpr, fwd_consts, *primals)
  return lu.wrap_init(fwd, debug_info=fwd_jaxpr_thunk.debug_info)

@lu.transformation2
def _handle_consts_in_bwd(f, const_avals, *args):
  cts, logs = f(*args)
  return [Zero(a) for a in const_avals] + list(cts), logs

custom_vjp_call_p = CustomVJPCallPrimitive('custom_vjp_call')
# TODO(phawkins,mattjj): make this primitive cacheable.
mlir.register_lowering(custom_vjp_call_p, _custom_jvp_vjp_call_lowering,
                       cacheable=False)

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 call_jaxpr.out_avals, core.positional_effects(call_jaxpr)
core.custom_typechecks[custom_vjp_call_p] = _custom_vjp_call_typecheck

def _custom_vjp_call_dce(
    used_outs: Sequence[bool], live_ins: Sequence[bool], eqn: core.JaxprEqn
) -> tuple[list[bool], core.JaxprEqn | None]:
  if not any(used_outs) and not pe.has_effects(eqn, live_ins):
    return [False] * len(eqn.invars), None
  call_jaxpr: core.Jaxpr = eqn.params["call_jaxpr"]
  fwd_jaxpr_thunk = eqn.params["fwd_jaxpr_thunk"]
  bwd: lu.WrappedFun = eqn.params["bwd"]
  out_trees: Callable[[], tuple[PyTreeDef, PyTreeDef, list[int | None]]] = eqn.params["out_trees"]
  dce_call_jaxpr: core.Jaxpr
  used_ins: Sequence[bool]
  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_, fwd_consts_ = fwd_jaxpr_thunk.call_wrapped(*zeros)
    fwd_jaxpr = fwd_jaxpr_.with_consts(fwd_consts_)
    _, res_tree, fwds = out_trees()
    num_res_out = res_tree.num_leaves - sum(f is not None for f in fwds)
    dce_fwd_jaxpr, _ = _cached_closed_call_dce_instantiate(
        fwd_jaxpr, (True,) * num_res_out + tuple(used_outs))
    return dce_fwd_jaxpr, dce_fwd_jaxpr.consts

  def dce_bwd(*args):
    _, res_tree, _ = out_trees()
    res, cts = split_list(args, [res_tree.num_leaves])
    cts_ = iter(cts)
    all_cts = []
    for used, aval in zip(used_outs, call_jaxpr.out_avals):
      if used:
        all_cts.append(next(cts_))
      else:
        all_cts.append(SymbolicZero(aval.to_ct_aval()))
    assert next(cts_, None) is None
    return bwd.call_wrapped(*res, *all_cts)

  dce_bwd_wrapped = lu.wrap_init(dce_bwd,
                                 debug_info=bwd.debug_info)
  outvars = [v for used, v in zip(used_outs, eqn.outvars) if used]
  new_params = dict(
      eqn.params,
      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,
      core.eqn_effects(dce_call_jaxpr, eqn.invars),
      eqn.source_info, eqn.ctx)
  return list(used_ins), new_eqn
pe.dce_rules[custom_vjp_call_p] = _custom_vjp_call_dce


def _custom_vjp_call_pp_rule(eqn: core.JaxprEqn,
                             context: core.JaxprPpContext,
                             settings: core.JaxprPpSettings) -> core.pp.Doc:
  params = dict(eqn.params)
  if not params["num_consts"]:
    params.pop("num_consts")
  params.pop("out_trees")
  params["fwd"] = params.pop("fwd_jaxpr_thunk").debug_info.func_name
  params["bwd"] = params.pop("bwd").debug_info.func_name
  names = sorted(params)
  params["name"] = params["call_jaxpr"].debug_info.func_name
  return core._pp_eqn(eqn.replace(params=params), context, settings,
                      params=["name"] + names)

core.pp_eqn_rules[custom_vjp_call_p] = _custom_vjp_call_pp_rule

batching.primitive_batchers[ad.custom_lin_p] = ad.raise_custom_vjp_error_on_jvp
# TODO(phawkins,mattjj): make this primitive cacheable.
mlir.register_lowering(ad.custom_lin_p, ad.raise_custom_vjp_error_on_jvp,
                       cacheable=False)


def custom_gradient(fun=None, *, with_logs: bool = False, remat: bool = False,
                    with_accums: Any = None):
  """Convenience function for defining custom VJP rules (aka custom gradients).

  While the canonical way to define custom VJP rules is via ``jax.custom_vjp``,
  the ``custom_gradient`` convenience wrapper follows TensorFlow's
  ``tf.custom_gradient`` API. The difference here is that ``custom_gradient``
  can be used as a decorator on one function that returns both the primal value
  (representing the output of the mathematical function to be differentiated)
  and the VJP (gradient) function. See
  https://www.tensorflow.org/api_docs/python/tf/custom_gradient.

  If the mathematical function to be differentiated has Haskell-like signature
  ``a -> b``, then the Python callable ``fun`` should have the signature
  ``a -> (b, CT b --o CT a)`` where we use ``CT x`` to denote a cotangent type
  for ``x`` and the ``--o`` arrow to denote a linear function. See the example
  below. That is, ``fun`` should return a pair where the first element
  represents the value of the mathematical function to be differentiated and the
  second element is a function to be called on the backward pass of reverse-mode
  automatic differentiation (i.e. the "custom gradient" function).

  The function returned as the second element of the output of ``fun`` can close
  over intermediate values computed when evaluating the function to be
  differentiated. That is, use lexical closure to share work between the forward
  pass and the backward pass of reverse-mode automatic differentiation. However,
  it cannot perform Python control flow which depends on the values of the
  closed-over intermediate values or its cotangent arguments; if the function
  includes such control flow, an error is raised.

  Args:
    fun: a Python callable specifying both the mathematical function to be
      differentiated and its reverse-mode differentiation rule. It should return
      a pair consisting of an output value and a Python callable that represents
      the custom gradient function.
    with_logs: optional bool, default ``False``. If ``True``, the custom
      gradient function must return a pair ``(in_cts, logs)`` rather than just
      the cotangents, where ``logs`` is a dict of named pytrees to log out of
      the backward pass, or ``None`` to log nothing, as with
      :py:meth:`jax.custom_vjp.defvjp_with_logs`.
      Log keys must be unique within a backward pass; repeated keys raise a
      ``ValueError``, even when the logs are ignored.
    remat: optional bool, default ``False``. If ``True``, ``fun`` also
      specifies custom rematerialization behavior, as with
      :py:func:`jax.custom_vjp.defremat`. The second element of the pair it
      returns is then a rematerialization function ``rem`` rather than a VJP
      function. ``rem`` takes the same arguments as ``fun`` and returns a
      pair of the output value and the VJP function. When differentiating
      under :func:`jax.remat`, ``rem`` runs on the backward pass, and the
      values it closes over are saved from the forward pass (so it should use
      its own arguments rather than closing over those of ``fun``); otherwise,
      it runs right after ``fun`` on the forward pass. Requires the
      ``jax_custom_vjp3`` and ``jax_remat3`` implementations.
    with_accums: optional, default ``None``. If given, the VJP function takes
      gradient accumulators rather than returning cotangents, as with
      :py:func:`jax.custom_vjp.defvjp_with_accums`: it's called as
      ``vjp(g, *arg_accums)``, with one accumulator per argument of ``fun``
      (matching its pytree structure), and returns ``None`` or a dict of
      backward-pass logs. The VJP function is traced on the forward pass, so
      ``with_accums`` declares which kinds of accumulators it gets: it's a
      pytree prefix of the tuple of arguments of ``fun`` (like the
      ``in_axes`` of :func:`jax.vmap`), with ``ValAccum`` or ``RefAccum``
      leaves. For example, ``(ValAccum, RefAccum)`` declares a ``ValAccum``
      for the first argument and a ``RefAccum`` for the second, and a single
      ``RefAccum`` declares a ``RefAccum`` for every argument. Inputs that
      aren't being differentiated, including any integer inputs, get a
      ``NullAccum`` whatever is declared. On the backward pass, the caller of
      autodiff must supply the declared kinds, except that it can discard the
      cotangent of an argument declared as a ``ValAccum`` (with
      ``jax.ad.DontWant()``); any other kind is an error. Can't yet be
      combined with ``remat``, or with the ``jax_custom_vjp3`` implementation.

  Returns:
    A Python callable that accepts the same arguments as ``fun`` and returns the
    output value specified by the first element of ``fun``'s output pair.

  For example:

  >>> @jax.custom_gradient
  ... def f(x):
  ...   return x ** 2, lambda g: (g * x,)
  ...
  >>> print(f(3.))
  9.0
  >>> print(jax.grad(f)(3.))
  3.0

  An example with a function on two arguments, so that the VJP function must
  return a tuple of length two:

  >>> @jax.custom_gradient
  ... def f(x, y):
  ...   return x * y, lambda g: (g * y, g * x)
  ...
  >>> print(f(3., 4.))
  12.0
  >>> print(jax.grad(f, argnums=(0, 1))(3., 4.))
  (Array(4., dtype=float32, weak_type=True), Array(3., dtype=float32, weak_type=True))

  With ``with_logs=True``, the VJP function returns a pair of the cotangents
  and a dict of backward-pass logs, received via the ``with_logs`` method of
  the VJP function that :py:func:`jax.vjp` returns:

  >>> @jax.custom_gradient(with_logs=True)
  ... def f(x):
  ...   return x ** 2, lambda g: ((g * 2 * x,), {'ct_out': g})
  ...
  >>> print(jax.grad(f)(3.))
  6.0
  >>> _, f_vjp = jax.vjp(f, 3.)
  >>> (x_ct,), logs = f_vjp.with_logs(1.)
  >>> print(logs['ct_out'])
  1.0

  With ``remat=True``, ``fun`` returns a rematerialization function in place of
  the VJP function. Here, under :func:`jax.remat`, the cosine is saved rather
  than recomputed on the backward pass::

    @jax.custom_gradient(remat=True)
    def sin(x):
      cos_x = jnp.cos(x)             # rem closes over it, so it's saved
      def rem(x):                    # runs on the backward pass
        return jnp.sin(x), lambda g: (g * cos_x,)
      return jnp.sin(x), rem
  """
  if fun is None:
    return lambda f: custom_gradient(f, with_logs=with_logs, remat=remat,
                                     with_accums=with_accums)
  if with_accums is not None and remat:
    raise NotImplementedError(
        "custom_gradient doesn't yet support with_accums together with remat")

  def wrapped_fun(*args, **kwargs):
    ans, _ = fun(*args, **kwargs)
    return ans

  wrapped_fun.__name__ = getattr(fun, '__name__', '<unnamed>')
  wrapped_fun.__qualname__ = getattr(fun, '__qualname__', '<unnamed>')
  wrapped_fun = custom_vjp(wrapped_fun)

  def fwd(*args, **kwargs):
    ans, rule = fun(*args, **kwargs)
    if with_logs:
      rule = _custom_gradient_logs_rule(rule)
    return ans, _closure_residuals("custom_gradient fwd", rule, (ans,), ct=True)

  def bwd(res, cts):
    in_tree = res.in_tree
    cts_out = _apply_closure_residuals(res, cts)
    if with_logs:
      cts_out, logs = cts_out
      cts_tree, _ = treedef_children(in_tree)
      if treedef_is_leaf(cts_tree):
        cts_out = (cts_out,)
      return cts_out, logs
    if treedef_is_leaf(in_tree):
      cts_out = (cts_out,)
    return cts_out

  if with_accums is not None:
    # symbolic_zeros tells us on the forward pass which inputs are perturbed,
    # and so will get NullAccums
    def accums_fwd(*primals):
      args = custom_vjp_primal_tree_values(primals)
      perturbed = tree_map(lambda p: p.perturbed, primals)
      ans, rule = fun(*args)
      kinds = _accum_kinds(with_accums, args, perturbed)
      return ans, _accums_closure_residuals(rule, ans, args, kinds)
    wrapped_fun.defvjp_with_accums(
        accums_fwd, partial(_apply_accums_closure_residuals, wrapped_fun.__name__),
        symbolic_zeros=True)
  elif remat:
    # The residuals saved on the forward pass are the values `rem` closes over,
    # and those for bwd are the values the VJP function closes over. Outside of
    # jax.remat, the VJP is the one defremat derives from these rules.
    def remat_fwd(*args):
      ans, rem = fun(*args)
      @wraps(rem)
      def rem_(*args):
        ans, rule = rem(*args)
        if with_logs:
          rule = _custom_gradient_logs_rule(rule)
        return ans, _closure_residuals("custom_gradient fwd", rule, (ans,), ct=True)
      return ans, _closure_residuals("custom_gradient rem", rem_, args, ct=False)
    defremat = wrapped_fun.defremat_with_logs if with_logs else wrapped_fun.defremat
    defremat(remat_fwd, _apply_closure_residuals, bwd)
  elif with_logs:
    wrapped_fun.defvjp_with_logs(fwd, bwd)
  else:
    wrapped_fun.defvjp(fwd, bwd)
  return wrapped_fun

def _closure_residuals(name, f, args, *, ct):
  """Trace `f` on the types (or cotangent types) of `args` into a `Residuals`,
  whose consts are the values `f` closes over."""
  args_flat, args_tree = tracing_registry.flatten((args, {}))
  avals = [core.typeof(x).to_ct_aval() if ct else core.typeof(x)
           for x in args_flat]
  closed_jaxpr, f_out = pe.trace_to_jaxpr(
      f, ft.treedef_args_to_ft(args_tree, avals), debug_info(name, f, args, {}))
  jaxpr, consts = pe.separate_consts(closed_jaxpr)
  return Residuals(jaxpr, f_out.tree, args_tree, consts)

def _apply_closure_residuals(res, *args):
  jaxpr, out_tree, args_tree, consts = res
  args_flat, args_tree_ = tree_flatten((args, {}))
  if args_tree != args_tree_: raise TypeError(f'{args_tree}\n!=\n{args_tree_}')
  return tree_unflatten(out_tree, core.eval_jaxpr(jaxpr, consts, *args_flat))

def _accum_kinds(spec, args, perturbed):
  """The kind of accumulator the VJP function gets for each leaf of `args`: a
  NullAccum if it isn't perturbed, and otherwise the kind `spec` declares."""
  kinds = broadcast_prefix(spec, args)
  for k in kinds:
    if k not in (ad.ValAccum, ad.RefAccum):
      raise TypeError(
          "custom_gradient's with_accums must be a pytree prefix of the "
          "function's arguments with ValAccum or RefAccum leaves, but got a "
          f"leaf {k}.")
  return tuple(k if p else ad.NullAccum
               for k, p in zip(kinds, tree_leaves(perturbed)))

def _accums_closure_residuals(rule, ans, args, kinds):
  """Trace `rule(g, *arg_accums)` with accumulators of the given kinds into a
  `Residuals`. The traced function returns logs, plus its ValAccums' values."""
  args_flat, args_tree = tree_flatten(args)
  ct_avals = [core.typeof(x).to_ct_aval() for x in args_flat]
  ref_avals = [AbstractRef(a) for a, k in zip(ct_avals, kinds) if k is ad.RefAccum]
  def rule_(g, refs):
    refs_ = iter(refs)
    accums = [ad.RefAccum(a, next(refs_)) if k is ad.RefAccum else k(a)
              for a, k in zip(ct_avals, kinds)]
    logs = rule(g, *tree_unflatten(args_tree, accums))
    vals = [acc.freeze() for acc in accums if isinstance(acc, ad.ValAccum)]
    return logs, [None if isinstance(v, Zero) else v for v in vals]
  ans_flat, _ = tracing_registry.flatten(ans)
  _, in_tree = tracing_registry.flatten(((ans, ref_avals), {}))
  avals = [*(core.typeof(x).to_ct_aval() for x in ans_flat), *ref_avals]
  closed_jaxpr, rule_out = pe.trace_to_jaxpr(
      rule_, ft.treedef_args_to_ft(in_tree, avals),
      debug_info("custom_gradient fwd", rule_, (ans, ref_avals), {}))
  jaxpr, consts = pe.separate_consts(closed_jaxpr)
  return Residuals(jaxpr, rule_out.tree, in_tree, consts, kinds)

def _apply_accums_closure_residuals(name, res, cts, *accums):
  jaxpr, out_tree, _, consts = res
  cts = tree_map(lambda c: zeros_like_aval(c.aval)
                 if isinstance(c, SymbolicZero) else c, cts)
  refs, val_accums = [], []
  for (path, acc), k in zip(tree_flatten_with_path(accums)[0], res.aux):
    if k is ad.ValAccum and isinstance(acc, (ad.ValAccum, ad.NullAccum)):
      val_accums.append(acc)  # a NullAccum here is from DontWant, and discards
    elif k is ad.RefAccum and isinstance(acc, ad.RefAccum):
      refs.append(acc.inst().ref)
    elif k is not ad.NullAccum:
      raise TypeError(
          f"custom_gradient function {name} declared a {k.__name__} for its "
          f"argument{keystr(path)} with with_accums, but its backward pass got "
          f"a {type(acc).__name__}. (A ValAccum comes from jax.grad, a VJP "
          "applied without with_refs, or transformations applied before "
          "differentiation, like vmap. A RefAccum comes from passing a Ref to "
          "a VJP's with_refs method, and a NullAccum from passing "
          "jax.ad.DontWant() to it.)")
  outs = core.eval_jaxpr(jaxpr, consts, *tree_leaves((cts, refs)))
  logs, vals = tree_unflatten(out_tree, outs)
  for acc, v in zip(val_accums, vals):
    if v is not None:
      acc.accum(v)
  return logs

def _custom_gradient_logs_rule(rule):
  @wraps(rule)
  def rule_with_logs(*cts):
    out = rule(*cts)
    if not (isinstance(out, (list, tuple)) and len(out) == 2):
      raise TypeError(
          "custom_gradient function used with with_logs=True must return a "
          "VJP function producing a pair (in_cts, logs), but the VJP function "
          f"returned {out}.")
    in_cts, logs = out
    if logs is not None and type(logs) is not dict:
      raise TypeError(
          "custom_gradient function used with with_logs=True must return a "
          "VJP function whose second output is None or a dict of "
          f"backward-pass log entries, but got {type(logs).__name__}.")
    return in_cts, logs
  return rule_with_logs

@register_pytree_node_class
class Residuals:
  def __init__(self, jaxpr, in_tree, out_tree, consts, aux=None):
    self.jaxpr = jaxpr
    self.in_tree = in_tree
    self.out_tree = out_tree
    self.consts = consts
    self.aux = aux  # other static data
  def __iter__(self):
    return iter((self.jaxpr, self.in_tree, self.out_tree, self.consts))
  def tree_flatten(self):
    return self.consts, (self.jaxpr, self.in_tree, self.out_tree, self.aux)
  @classmethod
  def tree_unflatten(cls, aux, consts):
    jaxpr, in_tree, out_tree, aux_ = aux
    return cls(jaxpr, in_tree, out_tree, consts, aux_)


def closure_convert(fun: Callable, *example_args) -> tuple[Callable, list[Any]]:
  """Closure conversion utility, for use with higher-order custom derivatives.

  To define custom derivatives such as with ``jax.custom_vjp(f)``, the target
  function ``f`` must take, as formal arguments, all values involved in
  differentiation. If ``f`` is a higher-order function, in that it accepts as an
  argument a Python function ``g``, then values stored away in ``g``'s closure
  will not be visible to the custom derivative rules, and attempts at AD
  involving these values will fail. One way around this is to convert the
  closure by extracting these values, and to pass them as explicit formal
  arguments across the custom derivative boundary. This utility carries out that
  conversion. More precisely, it closure-converts the function ``fun``
  specialized to the types of the arguments given in ``example_args``.

  When we refer here to "values in the closure" of ``fun``, we do not mean the
  values that are captured by Python directly when ``fun`` is defined (e.g. the
  Python objects in ``fun.__closure__``, if the attribute exists). Rather, we
  mean values encountered during the execution of ``fun`` on ``example_args``
  that determine its output. This may include, for instance, arrays captured
  transitively in Python closures, i.e. in the Python closure of functions
  called by ``fun``, the closures of the functions that they call, and so forth.

  The function ``fun`` must be a pure function.

  Example usage::

    def minimize(objective_fn, x0):
      converted_fn, aux_args = closure_convert(objective_fn, x0)
      return _minimize(converted_fn, x0, *aux_args)

    @partial(custom_vjp, nondiff_argnums=(0,))
    def _minimize(objective_fn, x0, *args):
      z = objective_fn(x0, *args)
      # ... find minimizer x_opt ...
      return x_opt

    def fwd(objective_fn, x0, *args):
      y = _minimize(objective_fn, x0, *args)
      return y, (y, args)

    def rev(objective_fn, res, g):
      y, args = res
      y_bar = g
      # ... custom reverse-mode AD ...
      return x0_bar, *args_bars

    _minimize.defvjp(fwd, rev)

  Args:
    fun: Python callable to be converted. Must be a pure function.
    example_args: Arrays, scalars, or (nested) standard Python
      containers (tuples, lists, dicts, namedtuples, i.e., pytrees)
      thereof, used to determine the types of the formal arguments to
      ``fun``. This type-specialized form of ``fun`` is the function
      that will be closure converted.

  Returns:
    A pair comprising (i) a Python callable, accepting the same
    arguments as ``fun`` followed by arguments corresponding to the
    values hoisted from its closure, and (ii) a list of values hoisted
    from the closure.
  """
  flat_args, in_tree = tracing_registry.flatten((example_args, {}))
  in_avals = tuple(map(core.typeof, flat_args))
  debug = debug_info("closure_convert", fun, example_args, {})
  if config.check_tracer_leaks.value:
    return _closure_convert_for_avals.__wrapped__(fun, in_tree, in_avals, debug)
  else:
    return _closure_convert_for_avals(fun, in_tree, in_avals, debug)

def _maybe_perturbed(x: Any) -> bool:
  # False if x can't represent an AD-perturbed value (i.e. a value
  # with a nontrivial tangent attached), up to heuristics, and True otherwise.
  # See https://github.com/jax-ml/jax/issues/6415 for motivation.
  if not isinstance(x, core.Tracer):
    # If x is not a Tracer, it can't be perturbed.
    return False
  elif isinstance(x, ad.JVPTracer) and isinstance(x.tangent, ad.Zero):
    return _maybe_perturbed(x.primal)
  elif isinstance(x, pe.DynamicJaxprTracer):
    # If x is a DynamicJaxprTracer then we're staging out; differentiation could
    # happen later, but some types always have trivial tangents.
    vspace = x.aval.to_tangent_aval()
    return not (vspace is core.abstract_token or
                getattr(vspace, 'dtype', None) == dtypes.float0)
  elif not isinstance(x, ad.JVPTracer):
    # If x is not a JVPTracer, recursively check its contents.
    return any(_maybe_perturbed(attr) for name, attr in x._contents())
  else:
    return True  # We can't be sure!

@cache()
def _closure_convert_for_avals(fun, in_tree, in_avals,
                               debug_info: core.DebugInfo):
  closed_jaxpr, out_avals = pe.trace_to_jaxpr(
      fun, ft.treedef_args_to_ft(in_tree, in_avals), debug_info)
  jaxpr, consts = pe.separate_consts(closed_jaxpr)
  out_tree = out_avals.tree

  (closure_consts, const_args), merge = partition_list(_maybe_perturbed, consts)
  num_consts = len(const_args)

  def converted_fun(*args_hconsts):
    num_args = len(args_hconsts) - num_consts
    args, const_args = split_list(args_hconsts, [num_args])
    consts = merge(closure_consts, const_args)
    all_args, in_tree2 = tree_flatten((tuple(args), {}))
    if in_tree != in_tree2:
      msg = ("The inputs to the closure produced by closure_convert must have "
             "the same Pytree structure as the example arguments passed when "
             f"closure_convert was called. Expected {in_tree}, but got "
             f"{in_tree2}")
      raise TypeError(msg)
    out_flat = core.eval_jaxpr(jaxpr, consts, *all_args)
    return tree_unflatten(out_tree, out_flat)

  return converted_fun, const_args

def partition_list(choice, lst):
  out = [], []
  which = [out[choice(elt)].append(elt) or choice(elt) for elt in lst]
  def merge(l1, l2):
    i1, i2 = iter(l1), iter(l2)
    return [next(i2 if snd else i1) for snd in which]
  return out, merge


### Custom transposition

def linear_call(fun: Callable,
                fun_transpose: Callable, residual_args,
                linear_args):
  """Call a linear function, with a custom implementation for its transpose.

  The `Haskell-like type signatures`_ of ``fun`` and ``fun_transpose`` are:

  .. code-block:: haskell

    fun           :: r -> a -o b
    fun_transpose :: r -> b -o a

  where the ``-o`` arrow indicates a linear function, ``r`` is the
  residual input type and ``a`` is the linear input type.

  The functions ``fun`` and ``fun_transpose`` are coupled as
  transposes of one another. Specifically, the transpose of a
  ``linear_call`` primitive is another ``linear_call`` to
  ``fun_transpose``, with ``fun`` as its custom transposition.

  For example:

  >>> def f(r, x):
  ...   return x / r

  >>> def t(r, t):
  ...   return t / r

  >>> def div_add(x, denom):
  ...   return x + linear_call(f, t, denom, x)

  >>> def transpose(f, x_example):
  ...   def transposed(y):
  ...     x, = jax.linear_transpose(f, x_example)(y)
  ...     return x
  ...   return transposed

  >>> div_add(9., 3.)
  Array(12., dtype=float32, weak_type=True)

  >>> transpose(partial(div_add, denom=3.), 1.)(18.)  # custom
  Array(24., dtype=float32, weak_type=True)

  >>> transpose(lambda x: x + x / 3., 1.)(18.)  # reference
  Array(24., dtype=float32, weak_type=True)

  The above definition of ``f`` illustrates the purpose of a residual
  argument: division is linear in one of its inputs (the dividend
  ``x``) but not the other (the divisor ``r``).

  As another example:

  >>> def custom_id(x):
  ...   def f(_, x): return x
  ...   def t(_, t): return 7.
  ...   return linear_call(f, t, (), x)
  >>> custom_id(1.)
  TypedFloat(1.0, dtype=float32)
  >>> transpose(custom_id, 1.)(1.)
  TypedFloat(7.0, dtype=float32)
  >>> transpose(transpose(custom_id, 1.), 1.)(1.)
  TypedFloat(1.0, dtype=float32)
  >>> transpose(transpose(transpose(custom_id, 1.), 1.), 1.)(1.)
  TypedFloat(7.0, dtype=float32)

  Args:
    fun: a Python callable specifying a linear function. It should
      take two arguments: one of "residual" inputs (type ``r``),
      i.e. inputs in which the function is not necessarily linear, and
      one of "linear" inputs (type ``a``).  It should return output
      whose components are linear in the linear input (type ``b``).
    fun_transpose: a Python callable specifying a structurally linear
      function that is the transpose of ``fun`` with respect to its
      linear inputs. Its first argument is the same residual inputs
      (``r``) as ``fun``. Its second argument is of type
      ``b``. Finally, its output is of type ``a`` and each of its
      component are linear in its second argument (the ``b`` inputs).
    residual_args: Argument in which ``fun`` and ``fun_transpose`` are
      not necessarily linear. Not involved in transposition.
    linear_args: Argument in which ``fun`` and ``fun_transpose`` are
      linear and with respect to which the two are transposes.

  Returns:
    The call result, i.e. ``fun(residual_args, linear_args)``.

  .. _Haskell-like type signatures: https://wiki.haskell.org/Type_signature
  """
  operands_res, res_tree = tracing_registry.flatten(residual_args)
  operands_lin, lin_tree = tracing_registry.flatten(linear_args)

  res_avals = map(core.typeof, operands_res)
  lin_avals = map(core.typeof, operands_lin)
  f_jaxpr, f_out_avals = pe.trace_to_jaxpr(
      fun,
      ft.pack(((ft.FTPyTree(res_avals, res_tree),
                ft.FTPyTree(lin_avals, lin_tree)), {})),
      debug_info("linear_call fun", fun, (residual_args, linear_args), {}))
  f_jaxpr_closed, f_consts = pe.separate_consts(f_jaxpr)
  out_avals = f_jaxpr_closed.out_avals
  out_tree = f_out_avals.tree

  @pe._memoize
  def transpose_thunk():
    t_jaxpr, t_out_avals = pe.trace_to_jaxpr(
        fun_transpose,
        ft.pack(((ft.FTPyTree(res_avals, res_tree),
                  ft.FTPyTree(list(out_avals), out_tree)), {})),
        # TODO(necula): the fun_transpose takes residual and output of fun!
        debug_info("linear_call fun_transpose", fun_transpose,
                   (residual_args, linear_args), {}).with_unknown_names())
    if t_out_avals.tree != lin_tree:
      raise TypeError(
          'transpose output pytree structure must match that of linear inputs, '
          f'got output structure {t_out_avals.tree} '
          f'and input structure {lin_tree}.')
    return pe.separate_consts(t_jaxpr)

  out = linear_call_p.bind(*f_consts, *operands_res, *operands_lin,
                           callee=f_jaxpr_closed,
                           transpose_thunk=transpose_thunk,
                           num_callee_consts=len(f_consts),
                           num_res=len(operands_res))

  return tree_unflatten(out_tree, out)

def _linear_call_impl(*args, callee, transpose_thunk, num_callee_consts,
                      num_res):
  del transpose_thunk, num_callee_consts, num_res
  return core.eval_jaxpr(callee, (), *args)

def _linear_call_jvp_rule(primals, tangents, callee, transpose_thunk,
                          num_callee_consts, num_res):
  consts_and_res, primals = split_list(primals, [num_callee_consts + num_res])
  const_tangents, tangents = split_list(tangents, [num_callee_consts + num_res])
  assert all(type(t) is Zero for t in const_tangents)
  primals_out = linear_call_p.bind(
      *consts_and_res, *primals, callee=callee, transpose_thunk=transpose_thunk,
      num_callee_consts=num_callee_consts, num_res=num_res)
  tangents_out = linear_call_p.bind(
      *consts_and_res, *tangents, callee=callee, transpose_thunk=transpose_thunk,
      num_callee_consts=num_callee_consts, num_res=num_res)
  return primals_out, tangents_out

def _linear_call_transpose_rule(cts, *args, callee, transpose_thunk,
                                num_callee_consts, num_res):
  transpose, t_consts = transpose_thunk()
  f_consts, operands_res, operands_lin = split_list(
      args, [num_callee_consts, num_res])
  _, _, cts_avals = split_list(
      transpose.in_avals, [len(t_consts), num_res])

  assert all(ad.is_undefined_primal(x)     for x in operands_lin)
  assert all(not ad.is_undefined_primal(x) for x in operands_res)

  def new_transpose_thunk():
    return callee, f_consts

  cts = [zeros_like_aval(a) if type(ct) is Zero else ct
         for ct, a in zip(cts, cts_avals)]
  cts_out = linear_call_p.bind(*t_consts, *operands_res, *cts,
                               callee=transpose,
                               transpose_thunk=new_transpose_thunk,
                               num_callee_consts=len(t_consts),
                               num_res=len(operands_res))

  return [None] * (num_callee_consts + num_res) + cts_out

def _linear_call_abstract_eval(*args, **kwargs):
  return kwargs['callee'].out_avals

linear_call_p = core.Primitive('linear_call')
linear_call_p.multiple_results = True
linear_call_p.def_impl(_linear_call_impl)
linear_call_p.def_abstract_eval(_linear_call_abstract_eval)
ad.primitive_jvps[linear_call_p] = _linear_call_jvp_rule
ad.primitive_transposes[linear_call_p] = _linear_call_transpose_rule
mlir.register_lowering(linear_call_p, mlir.lower_fun(
    _linear_call_impl, multiple_results=True))


# A stageable primitive that fails when evaluated
unreachable_p: core.Primitive = core.Primitive('unreachable')
unreachable_p.multiple_results = True

def unreachable_impl(*_, out_avals, exc_type, message):
  del out_avals
  raise exc_type(message)

# Evaluation raises an exception
unreachable_p.def_impl(unreachable_impl)

# Translation raises an exception
# TODO(frostig,mattjj): We have no good way to translate a function
# that errs. Since MLIR lowering over-approximates concrete evaluation,
# we err on MLIR lowering for the time being.
mlir.register_lowering(unreachable_p, unreachable_impl)

# Abstract evaluation proceeds without issue, to allow for staging
unreachable_p.def_abstract_eval(lambda *_, out_avals, **__: out_avals)

def unreachable(*args, out_avals=None, exc_type=TypeError,
                message='unreachable'):
  """Fail when evaluated concretely (but allow for staging).

  This function allows one to assert an impossibility of
  evaluation. It can be used to guarantee that evaluation does not
  "reach" a certain point in the sense that it does not execute, but
  it can nonetheless be staged out by JAX without error.

  Args:
    *args: The arbitrary pytree of arguments to the function.
    out_avals: Optional specification of the output types of this
     function invocation from the point of view of staging. If
     ``None``, these are chosen as equal to types of input arguments.
    exc_type: Optional constructor for the Python exception raised if
      evaluated.
    message: Optional string message for the Python exception raised
      if evaluated.

  """
  if out_avals is None:
    out_avals = tree_map(core.typeof, args)

  args_flat, in_tree = tree_flatten(args)
  out_avals_flat, out_tree = tree_flatten(out_avals)
  out = unreachable_p.bind(*args_flat, out_avals=out_avals_flat,
                           exc_type=exc_type, message=message)
  return tree_unflatten(out_tree, out)


disallow_jvp = partial(
    unreachable,
    exc_type=TypeError,
    message="can't apply forward-mode autodiff (jvp) to a custom_vjp function.")


# 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
# automatically.
#
# TODO(dfm): Eventually this should probably be the default behavior for
# custom_vjp, if we can make it so that it is a no-op for most cases. Right now,
# it is written in "initial-style" so it doesn't support eager mode. This was
# a reasonable compromise when written because it made the implementation
# simpler, but it would be worth revisiting this.
def optimize_remat_of_custom_vjp_fwd[ReturnValue](
    fun: Callable[..., ReturnValue],
    debug_fun: core.DebugInfo,
    fwd: Callable[..., tuple[ReturnValue, Any]],
    debug_fwd: core.DebugInfo,
    nondiff_argnums: Sequence[int] = (),
    symbolic_zeros: bool = False,
) -> Callable[..., tuple[ReturnValue, Any]]:
  if symbolic_zeros:
    # TODO(dfm): This probably shouldn't be too hard to support.
    raise NotImplementedError(
        "remat optimization for custom_vjp does not support symbolic zeros")

  @wraps(fwd)
  def wrapped_fwd(*args, **kwargs) -> tuple[ReturnValue, Any]:
    # TODO(dfm): This initial logic is duplicated from custom_vjp.__call__
    # above and it would be good to consolidate it.
    # Note: we use `fun` instead of `fwd` here for consistency with
    # custom_vjp.__call__ above.
    args = resolve_kwargs(fun, args, kwargs)
    if nondiff_argnums:
      for i in nondiff_argnums: _check_for_tracers(args[i])
      nondiff_argnums_ = set(nondiff_argnums)
      dyn_argnums = [i for i in range(len(args)) if i not in nondiff_argnums_]
      f_, dyn_args = argnums_partial(lu.wrap_init(fun, debug_info=debug_fun),
                                     dyn_argnums,
                                     args, require_static_args_hashable=False)
      fwd_, _ = argnums_partial(lu.wrap_init(fwd, debug_info=debug_fwd),
                                dyn_argnums, args,
                                require_static_args_hashable=False)
    else:
      f_, dyn_args = lu.wrap_init(fun, debug_info=debug_fun), args
      fwd_ = lu.wrap_init(fwd, debug_info=debug_fwd)
    args_flat, in_tree = tree_flatten(dyn_args)
    flat_fun, out_type = _flatten_fun_nokwargs(f_, in_tree)
    flat_fwd, out_trees = _flatten_fwd(fwd_, False, debug_fun, debug_fwd,
                                       in_tree, out_type)
    flat_fwd = _fix_fwd_args(flat_fwd)

    in_avals = [core.typeof(x) for x in args_flat]
    fwd_jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(flat_fwd.with_unknown_names(),
                                                     in_avals)
    fwd_jaxpr = pe.convert_constvars_jaxpr(fwd_jaxpr)
    prim_tree, res_tree, fwds = out_trees()
    num_res_out = res_tree.num_leaves - sum(f is not None for f in fwds)

    disallowed_effects = effects.custom_derivatives_allowed_effects.filter_not_in(fwd_jaxpr.effects)
    if disallowed_effects:
      raise NotImplementedError(
          "remat optimization for custom_vjp does not support forward "
          f"functions with these side effects: {disallowed_effects}")

    @pe._memoize
    def fun_jaxpr_thunk():
      jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(flat_fun, in_avals)
      return jaxpr, consts

    out_flat = remat_opt_p.bind(*consts, *args_flat, num_consts=len(consts),
                                num_res=num_res_out, fwd_jaxpr=fwd_jaxpr,
                                fun_jaxpr_thunk=fun_jaxpr_thunk)
    res, out_flat = split_list(out_flat, [num_res_out])
    res_ = iter(res)
    res = [next(res_) if f is None else args_flat[f] for f in fwds]
    assert next(res_, None) is None
    out_tree = treedef_tuple((prim_tree, res_tree))
    return tree_unflatten(out_tree, (*out_flat, *res))

  return wrapped_fwd

@lu.transformation2
def _fix_fwd_args(f, *args):
  args = [(x, True) for x in args]
  args = [x for pair in args for x in pair]
  return f(*args)

def _remat_opt_impl(
    *args,
    num_consts: int,
    num_res: int,
    fwd_jaxpr: core.Jaxpr,
    fun_jaxpr_thunk: Callable[[], tuple[core.Jaxpr, Sequence[Any]]],
):
  del num_consts, num_res, fun_jaxpr_thunk  # unused
  return core.jaxpr_as_fun(fwd_jaxpr)(*args)

def _remat_opt_abstract_eval(*args, fwd_jaxpr: core.Jaxpr, **_):
  del args
  return fwd_jaxpr.out_avals, core.positional_effects(fwd_jaxpr)

def _remat_opt_vmap(
    axis_data, args, in_dims,
    *,
    num_consts: int,
    num_res: int,
    fwd_jaxpr: core.Jaxpr,
    fun_jaxpr_thunk: Callable[[], tuple[core.Jaxpr, Sequence[Any]]],
):
  args = [batching.moveaxis(x, d, 0) if d is not None and d != 0
          else x for x, d in zip(args, in_dims)]
  in_batched = [d is not None for d in in_dims]
  batched_fwd_jaxpr, out_batched = batching.batch_jaxpr(
      fwd_jaxpr, axis_data, in_batched, False)
  extra_consts = batched_fwd_jaxpr.consts
  batched_fwd_jaxpr = pe.convert_constvars_jaxpr(batched_fwd_jaxpr)
  out_dims = [0 if b else None for b in out_batched]

  _, prim_batched = split_list(in_batched, [num_consts])

  @pe._memoize
  def batched_fun_jaxpr_thunk():
    fun_jaxpr_, fun_consts_ = fun_jaxpr_thunk()
    fun_jaxpr = fun_jaxpr_.with_consts(fun_consts_)
    batched_fun_jaxpr, out_batched = batching.batch_jaxpr(
        fun_jaxpr, axis_data, prim_batched, False)
    return batched_fun_jaxpr, batched_fun_jaxpr.consts

  batched_outs = remat_opt_p.bind(*extra_consts, *args,
                                  num_consts=num_consts + len(extra_consts),
                                  num_res=num_res,
                                  fwd_jaxpr=batched_fwd_jaxpr,
                                  fun_jaxpr_thunk=batched_fun_jaxpr_thunk)

  return batched_outs, out_dims

def _remat_opt_jvp(
    primals,
    tangents,
    *,
    num_consts: int,
    num_res: int,
    fwd_jaxpr: core.Jaxpr,
    fun_jaxpr_thunk: Callable[[], tuple[core.Jaxpr, Sequence[Any]]],
):
  consts, primals = split_list(primals, [num_consts])
  consts_dot, tangents = split_list(tangents, [num_consts])
  # Tangents must be instantated in case we end up DCEing later.
  tangents = map(ad.instantiate_zeros, tangents)
  consts_nz = [not isinstance(t, Zero) for t in consts_dot]
  consts_dot = [c for nz, c in zip(consts_nz, consts_dot) if nz]
  in_nz = consts_nz + [True] * len(tangents)
  fwd_jaxpr_jvp_, out_nz = ad.jvp_jaxpr(fwd_jaxpr, in_nz, True)
  num_out = len(out_nz) - num_res
  fwd_jaxpr_jvp_ = ad.rearrange_binders(
      fwd_jaxpr_jvp_, [num_consts, len(primals)],
      [len(consts_dot), len(tangents)], [num_res, num_out], [num_res, num_out])
  fwd_jaxpr_jvp = pe.convert_constvars_jaxpr(fwd_jaxpr_jvp_)

  # @pe._memoize
  def fun_jvp_jaxpr_thunk():
    fun_jaxpr_, fun_consts_ = fun_jaxpr_thunk()
    fun_jaxpr = fun_jaxpr_.with_consts(fun_consts_)
    in_nz = [True] * len(primals)
    fun_jvp_jaxpr, _ = ad.jvp_jaxpr(fun_jaxpr, in_nz, True)
    return fun_jvp_jaxpr, fun_jvp_jaxpr.consts

  new_num_consts = len(fwd_jaxpr_jvp_.consts) + num_consts + len(consts_dot)
  outs = remat_opt_p.bind(*fwd_jaxpr_jvp_.consts, *consts, *consts_dot,
                          *primals, *tangents, num_consts=new_num_consts,
                          num_res=2 * num_res, fwd_jaxpr=fwd_jaxpr_jvp,
                          fun_jaxpr_thunk=fun_jvp_jaxpr_thunk)
  res, res_dot, outs, outs_dot = split_list(outs, [num_res, num_res, num_out])
  return (*res, *outs), (*res_dot, *outs_dot)

def _remat_opt_transpose(
    cts, *args,
    num_consts: int,
    num_res: int,
    fwd_jaxpr: core.Jaxpr,
    fun_jaxpr_thunk: Callable[[], tuple[core.Jaxpr, Sequence[Any]]],
):
  # TODO(dfm): It shouldn't be too hard to implement this as needed in the
  # future.
  raise NotImplementedError(
      "remat optimization for custom_vjp does not support higher-order AD")

def _remat_opt_dce(used_outs: list[bool], live_ins: list[bool],
                   eqn: core.JaxprEqn):
  if not any(used_outs) and not pe.has_effects(eqn, live_ins):
    return [False] * len(eqn.invars), None
  used_res, used_prims = split_list(used_outs, [eqn.params["num_res"]])
  outvars = [v for used, v in zip(used_outs, eqn.outvars) if used]
  if any(used_res):
    # If any of the residuals are used, we still need to run fwd at this point,
    # but we may end up DCEing again in the future, so we must instantiate all
    # the input primals.
    instantiate = [False] * eqn.params["num_consts"]
    instantiate += [True] * (len(eqn.invars) - eqn.params["num_consts"])
    new_jaxpr, used_ins = pe.dce_jaxpr(eqn.params["fwd_jaxpr"], used_outs,
                                       instantiate=instantiate)
    assert not new_jaxpr.constvars
    closed_jaxpr = new_jaxpr
    invars = [v for used, v in zip(used_ins, eqn.invars) if used]
    new_params = dict(eqn.params)
    new_num_consts = sum(split_list(used_ins, [eqn.params["num_consts"]])[0])
    new_params["num_consts"] = new_num_consts
    new_params["fwd_jaxpr"] = closed_jaxpr
    new_params["num_res"] = sum(used_res)
    new_eqn = pe.new_jaxpr_eqn(
        invars, outvars, remat_opt_p, new_params,
        core.eqn_effects(closed_jaxpr, invars),
        eqn.source_info, eqn.ctx)
    return used_ins, new_eqn
  else:
    # If none of the residuals are used, we run the primal computation instead.
    # At this point we drop this custom DCE behavior, but since the primal might
    # have different consts than fwd, we build a new JaxprEqn with a closed_call
    # primitive.
    fun_jaxpr, consts = eqn.params["fun_jaxpr_thunk"]()
    closed_jaxpr, _, used_ins = pe.dce_jaxpr_consts(
        fun_jaxpr.with_consts(consts), used_prims)
    _, invars = split_list(eqn.invars, [eqn.params["num_consts"]])
    invars = [v for used, v in zip(used_ins, invars) if used]
    new_eqn = pe.new_jaxpr_eqn(
        invars, outvars, core.eval_jaxpr_p, dict(call_jaxpr=closed_jaxpr),
        core.eqn_effects(closed_jaxpr, invars), eqn.source_info, eqn.ctx)
    used_ins = [False] * eqn.params["num_consts"] + used_ins
    return used_ins, new_eqn

def _remat_opt_to_lojax(*hi_args, fwd_jaxpr: core.Jaxpr, num_consts, **params):
  return pe._lower_and_eval(pe.eval_jaxpr_p, fwd_jaxpr, hi_args)

remat_opt_p = core.Primitive("remat_opt")
remat_opt_p.multiple_results = True
remat_opt_p.is_high = lambda *_, fwd_jaxpr, **__: fwd_jaxpr.is_high
remat_opt_p.to_lojax = _remat_opt_to_lojax
remat_opt_p.def_impl(_remat_opt_impl)
remat_opt_p.def_effectful_abstract_eval(_remat_opt_abstract_eval)
mlir.register_lowering(remat_opt_p, mlir.lower_fun(
    _remat_opt_impl, multiple_results=True))


batching.fancy_primitive_batchers[remat_opt_p] = _remat_opt_vmap
ad.primitive_jvps[remat_opt_p] = _remat_opt_jvp
ad.primitive_transposes[remat_opt_p] = _remat_opt_transpose
pe.dce_rules[remat_opt_p] = _remat_opt_dce