{"record":{"id":"3a29f8cb7cce94b8","repo":"jax-ml/jax","slug":"remat-optimization-for-custom-vjp-does-not-support-3a29f8","errorCode":null,"errorMessage":"remat optimization for custom_vjp does not support forward functions with these side effects: {disallowed_effects}","messagePattern":"remat optimization for custom_vjp does not support forward functions with these side effects: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_derivatives.py","lineNumber":1761,"sourceCode":"    else:\n      f_, dyn_args = lu.wrap_init(fun, debug_info=debug_fun), args\n      fwd_ = lu.wrap_init(fwd, debug_info=debug_fwd)\n    args_flat, in_tree = tree_flatten(dyn_args)\n    flat_fun, out_type = _flatten_fun_nokwargs(f_, in_tree)\n    flat_fwd, out_trees = _flatten_fwd(fwd_, nondiff_argnums, False,\n                                       debug_fun, debug_fwd, in_tree, out_type)\n    flat_fwd = _fix_fwd_args(flat_fwd)\n\n    in_avals = [core.typeof(x) for x in args_flat]\n    fwd_jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(flat_fwd.with_unknown_names(),\n                                                     in_avals)\n    fwd_jaxpr = pe.convert_constvars_jaxpr(fwd_jaxpr)\n    prim_tree, res_tree, fwds = out_trees()\n    num_res_out = res_tree.num_leaves - sum(f is not None for f in fwds)\n\n    disallowed_effects = effects.custom_derivatives_allowed_effects.filter_not_in(fwd_jaxpr.effects)\n    if disallowed_effects:\n      raise NotImplementedError(\n          \"remat optimization for custom_vjp does not support forward \"\n          f\"functions with these side effects: {disallowed_effects}\")\n\n    @pe._memoize\n    def fun_jaxpr_thunk():\n      jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(flat_fun, in_avals)\n      return jaxpr, consts\n\n    out_flat = remat_opt_p.bind(*consts, *args_flat, num_consts=len(consts),\n                                num_res=num_res_out, fwd_jaxpr=fwd_jaxpr,\n                                fun_jaxpr_thunk=fun_jaxpr_thunk)\n    res, out_flat = split_list(out_flat, [num_res_out])\n    res_ = iter(res)\n    res = [next(res_) if f is None else args_flat[f] for f in fwds]\n    assert next(res_, None) is None\n    out_tree = treedef_tuple((prim_tree, res_tree))\n    return tree_unflatten(out_tree, (*out_flat, *res))\n","sourceCodeStart":1743,"sourceCodeEnd":1779,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_derivatives.py#L1743-L1779","documentation":"The remat optimization for jax.custom_vjp only allows effects in the effects.custom_derivatives_allowed_effects set. If the traced forward function performs side effects outside that allowlist (e.g. ordered/ordered-in-effect prints, random state updates, or custom effects), NotImplementedError is raised.","triggerScenarios":"Applying remat-optimized custom_vjp to a forward function that prints with ordered effects, uses stateful primitives, or declares custom effects not permitted for custom derivatives.","commonSituations":"Debug prints (jax.debug.print with ordered=True) or PRNG handling inside the custom_vjp forward function combined with checkpointing.","solutions":["Remove disallowed side effects (ordered prints, stateful ops) from the custom_vjp forward function","Move prints/effects outside the checkpointed region","Fall back to plain remat without the custom_vjp remat optimization (e.g. jax.checkpoint on an outer function)"],"exampleFix":"# before\n@custom_vjp\ndef f(x):\n  jax.debug.print(\"x={}\", x, ordered=True)  # disallowed effect\n  ...\n\n# after\n@custom_vjp\ndef f(x):\n  ...  # no side effects inside","handlingStrategy":"validation","validationCode":"jaxpr = jax.make_jaxpr(fwd)(x)\nallowed = jax._src.effects.effects.custom_derivatives_allowed_effects\nassert not jaxpr.effects - allowed, jaxpr.effects - allowed","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep custom_vjp forward functions pure: no ordered prints, no stateful ops","Check jaxpr.effects during development of custom ops"],"tags":["jax","custom-vjp","remat","side-effects"],"backgroundTag":"disallowed-side-effects-in-traced-code","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}