jax-ml/jax · error · NotImplementedError

remat optimization for custom_vjp does not support higher-or

Error message

remat optimization for custom_vjp does not support higher-order AD

What it means

jax.remat applied to a custom_vjp produces a Jaxpr whose transposition is not implemented for higher-order automatic differentiation. Taking a second (or higher) derivative through a remat-optimized custom_vjp hits this NotImplementedError.

Source

Thrown at jax/_src/custom_derivatives.py:1885

  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], eqn: core.JaxprEqn):
  if not any(used_outs) and not pe.has_effects(eqn):
    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]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove remat from the custom_vjp function when computing higher-order derivatives
2. Implement the custom op with custom_jvp (which supports higher-order AD) instead of custom_vjp
3. Checkpoint a different (non-custom_vjp) part of the computation

Example fix

# before
loss = jax.checkpoint(custom_vjp_fn)(x)
h = jax.grad(jax.grad(loss))(x)  # higher-order AD -> error

# after
custom_vjp_fn_noremat = custom_vjp_fn  # without checkpoint
h = jax.grad(jax.grad(custom_vjp_fn_noremat))(x)

try:
    h = jax.grad(jax.grad(loss))(x)
except NotImplementedError as e:
    if 'higher-order AD' in str(e):
        h = jax.grad(jax.grad(loss_plain))(x)  # without remat/custom_vjp

When it happens

Trigger: Calling grad twice (or vjp of grad, jacfwd of grad, etc.) on a function containing remat(custom_vjp function).

Common situations: Computing Hessians or gradient penalties of loss functions that use custom_vjp (e.g. custom simplices, trilinear interpolation losses) with checkpointing for memory.

Related errors

• remat optimization for custom_vjp does not support forward f
• Pure callbacks do not support transpose. Please use `jax.cus
• under vmap, the {member_name} of {name} produced an output b
• Custom VJP bwd rule {self.bwd} was registered with defvjp_wi
• Custom VJP bwd rule {self.bwd} was registered with defvjp_wi

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/8d902cb9abe58737. Report an issue: GitHub.