jax-ml/jax · error · NotImplementedError

remat optimization for custom_vjp does not support forward f

Error message

remat optimization for custom_vjp does not support forward functions with these side effects: {disallowed_effects}

What it means

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.

Source

Thrown at jax/_src/custom_derivatives.py:1761

    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_, nondiff_argnums, 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))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove disallowed side effects (ordered prints, stateful ops) from the custom_vjp forward function
2. Move prints/effects outside the checkpointed region
3. Fall back to plain remat without the custom_vjp remat optimization (e.g. jax.checkpoint on an outer function)

Example fix

# before
@custom_vjp
def f(x):
  jax.debug.print("x={}", x, ordered=True)  # disallowed effect
  ...

# after
@custom_vjp
def f(x):
  ...  # no side effects inside

jaxpr = jax.make_jaxpr(fwd)(x)
allowed = jax._src.effects.effects.custom_derivatives_allowed_effects
assert not jaxpr.effects - allowed, jaxpr.effects - allowed

When it happens

Trigger: 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.

Common situations: Debug prints (jax.debug.print with ordered=True) or PRNG handling inside the custom_vjp forward function combined with checkpointing.

Related errors

• remat optimization for custom_vjp does not support higher-or
• 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/3a29f8cb7cce94b8. Report an issue: GitHub.