jax-ml/jax · error · NotImplementedError

Effects not supported in `scan`: {disallowed_effects}

Error message

Effects not supported in `scan`: {disallowed_effects}

What it means

The new (scan3) implementation of lax.scan rejects jaxprs containing effects outside the control-flow whitelist. If the scan body performs IO or other disallowed effects, tracing fails with NotImplementedError.

Source

Thrown at jax/_src/lax/control_flow/loops.py:190

  xs_avals = xs_flat.map(core.typeof)
  length = _infer_scan_length(list(xs_flat), list(xs_avals), length)

  # TODO(dougalm): handle disable_jit
  if config.mutable_array_checks.value:
    check_no_aliased_ref_args(lambda: dbg_body, list(xs_avals), list(xs_flat))

  x_avals = xs_avals.map(lambda aval: core.mapped_leading_aval(length, aval))
  # TODO(dougalm): promote away all weak types
  args_avals = ft.pack(((x_avals,), {}))
  jaxpr, y_avals = pe.trace_to_jaxpr(f, args_avals, dbg_body)
  jaxpr, consts = pe.separate_consts(jaxpr)

  if config.mutable_array_checks.value:
    _check_no_aliased_closed_over_refs(dbg_body, consts, list(xs_flat))

  disallowed_effects = effects.control_flow_allowed_effects.filter_not_in(jaxpr.effects)
  if disallowed_effects:
    raise NotImplementedError(
        f'Effects not supported in `scan`: {disallowed_effects}')

  unroll = core.concrete_or_error(
      None, unroll,
      "The `unroll` argument to `scan` expects a concrete `int` or `bool` "
      "value.")
  if isinstance(unroll, bool):
    unroll = max(length, 1) if unroll else 1
  if unroll < 0:
    raise ValueError("`unroll` must be a `bool` or a non-negative `int`.")

  args = list(consts) + list(xs_flat)
  # TODO(dougalm): handle traceable-level forwarding
  out = Scan3(
      extensives = [False] * len(consts) + [True] * len(xs_flat),
      length=length, jaxpr=jaxpr, reverse=reverse, unroll=unroll)(args)

  return y_avals.update(out).unflatten()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove effectful ops from the scan body; hoist IO outside the scan
2. Compute what the callback needs purely and do IO on the final result
3. Disable scan3 via jax.config.update('jax_scan3', False) if you need the old behavior temporarily
4. Register your custom effect in control_flow_allowed_effects if you own it

Example fix

// before
jax.config.update('jax_scan3', True)
def body(c, x): hcb.io_callback(print, None, x); return c, x
// after
def body(c, x): return c, x  # IO removed; print after scan on results

j = jax.make_jaxpr(body)(carry_example, x_example)
from jax._src.effects import control_flow_allowed_effects
assert not control_flow_allowed_effects.filter_not_in(j.effects)

def scan_body_is_pure(body, c_ex, x_ex) -> bool:
    return not control_flow_allowed_effects.filter_not_in(jax.make_jaxpr(body)(c_ex, x_ex).effects)

try: jax.config.update('jax_scan3', True); lax.scan(...)\nexcept NotImplementedError as e:\n    if 'Effects not supported' in str(e): jax.config.update('jax_scan3', False); retry
    else: raise

When it happens

Trigger: Using jax.config.update('jax_scan3', True) (or a version where scan3 is default) with a scan body that calls io_callback, debug printers, or custom effectful primitives.

Common situations: Opting into the experimental scan3 path while body still contains callbacks; version upgrades where scan3 became default and previously tolerated effects no longer pass.

Related errors

• must provide `length` to `scan`, since the leading-axis size
• Effects not supported in `custom_jvp`: {disallowed_effects}
• Effects not supported in `custom_jvp`: {disallowed}
• Cannot lower jaxpr with effects: {closed_jaxpr.effects}
• `JaxprInputEffect` {eff} does not have a corresponding jaxpr

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