{"record":{"id":"8d902cb9abe58737","repo":"jax-ml/jax","slug":"remat-optimization-for-custom-vjp-does-not-support-8d902c","errorCode":null,"errorMessage":"remat optimization for custom_vjp does not support higher-order AD","messagePattern":"remat optimization for custom_vjp does not support higher-order AD","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_derivatives.py","lineNumber":1885,"sourceCode":"\n  new_num_consts = len(fwd_jaxpr_jvp_.consts) + num_consts + len(consts_dot)\n  outs = remat_opt_p.bind(*fwd_jaxpr_jvp_.consts, *consts, *consts_dot,\n                          *primals, *tangents, num_consts=new_num_consts,\n                          num_res=2 * num_res, fwd_jaxpr=fwd_jaxpr_jvp,\n                          fun_jaxpr_thunk=fun_jvp_jaxpr_thunk)\n  res, res_dot, outs, outs_dot = split_list(outs, [num_res, num_res, num_out])\n  return (*res, *outs), (*res_dot, *outs_dot)\n\ndef _remat_opt_transpose(\n    cts, *args,\n    num_consts: int,\n    num_res: int,\n    fwd_jaxpr: core.Jaxpr,\n    fun_jaxpr_thunk: Callable[[], tuple[core.Jaxpr, Sequence[Any]]],\n):\n  # TODO(dfm): It shouldn't be too hard to implement this as needed in the\n  # future.\n  raise NotImplementedError(\n      \"remat optimization for custom_vjp does not support higher-order AD\")\n\ndef _remat_opt_dce(used_outs: list[bool], eqn: core.JaxprEqn):\n  if not any(used_outs) and not pe.has_effects(eqn):\n    return [False] * len(eqn.invars), None\n  used_res, used_prims = split_list(used_outs, [eqn.params[\"num_res\"]])\n  outvars = [v for used, v in zip(used_outs, eqn.outvars) if used]\n  if any(used_res):\n    # If any of the residuals are used, we still need to run fwd at this point,\n    # but we may end up DCEing again in the future, so we must instantiate all\n    # the input primals.\n    instantiate = [False] * eqn.params[\"num_consts\"]\n    instantiate += [True] * (len(eqn.invars) - eqn.params[\"num_consts\"])\n    new_jaxpr, used_ins = pe.dce_jaxpr(eqn.params[\"fwd_jaxpr\"], used_outs,\n                                       instantiate=instantiate)\n    assert not new_jaxpr.constvars\n    closed_jaxpr = new_jaxpr\n    invars = [v for used, v in zip(used_ins, eqn.invars) if used]","sourceCodeStart":1867,"sourceCodeEnd":1903,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_derivatives.py#L1867-L1903","documentation":"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.","triggerScenarios":"Calling grad twice (or vjp of grad, jacfwd of grad, etc.) on a function containing remat(custom_vjp function).","commonSituations":"Computing Hessians or gradient penalties of loss functions that use custom_vjp (e.g. custom simplices, trilinear interpolation losses) with checkpointing for memory.","solutions":["Remove remat from the custom_vjp function when computing higher-order derivatives","Implement the custom op with custom_jvp (which supports higher-order AD) instead of custom_vjp","Checkpoint a different (non-custom_vjp) part of the computation"],"exampleFix":"# before\nloss = jax.checkpoint(custom_vjp_fn)(x)\nh = jax.grad(jax.grad(loss))(x)  # higher-order AD -> error\n\n# after\ncustom_vjp_fn_noremat = custom_vjp_fn  # without checkpoint\nh = jax.grad(jax.grad(custom_vjp_fn_noremat))(x)","handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":"try:\n    h = jax.grad(jax.grad(loss))(x)\nexcept NotImplementedError as e:\n    if 'higher-order AD' in str(e):\n        h = jax.grad(jax.grad(loss_plain))(x)  # without remat/custom_vjp","preventionTips":["Keep a non-checkpointed variant of custom_vjp functions for Hessian computations","Test second derivatives early when designing custom ops"],"tags":["jax","custom-vjp","remat","higher-order-ad"],"backgroundTag":"higher-order-ad-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}