jax-ml/jax · error · ValueError

The following ordered effects are not supported for more tha

Error message

The following ordered effects are not supported for more than 1 device: {unsupported_effects}

What it means

Ordered effects (e.g. host callback ordered effects like print with ordered=True, or custom ordered effects) execute in sequence, which cannot be preserved when a computation is sharded across multiple devices. Lowering (via _cached_lowering_to_hlo) checks this and rejects ordered, non-shardable effects whenever num_devices > 1.

Source

Thrown at jax/_src/interpreters/pxla.py:748

  log_priority = logging.WARNING if config.log_compiles.value else logging.DEBUG
  if logger.isEnabledFor(log_priority):
    logger.log(log_priority,
               "Compiling %s with global shapes and types %s. "
               "Argument mapping: %s.",
               module_name, in_avals, in_shardings)

  in_mlir_shardings = map(_to_logical_sharding, in_avals, in_shardings)
  out_mlir_shardings = map(_to_logical_sharding, out_avals, out_shardings)
  replicated_args = [False] * len(in_avals)
  axis_ctx = sharding_impls.ShardingContext(num_devices, device_assignment,
                                            abstract_mesh)

  if num_devices > 1:
    unsupported_effects = effects.ordered_effects.filter_in(closed_jaxpr.effects)
    unsupported_effects = effects.shardable_ordered_effects.filter_not_in(
        unsupported_effects)
    if len(unsupported_effects) > 0:
      raise ValueError(
        "The following ordered effects are not supported for "
        f"more than 1 device: {unsupported_effects}")
  ordered_effects = list(effects.ordered_effects.filter_in(closed_jaxpr.effects))
  arg_names = ("",) * num_const_args + jaxpr._debug_info.safe_arg_names(len(in_avals) - num_const_args)
  with dispatch.log_elapsed_time(
        "Finished jaxpr to MLIR module conversion {fun_name} in {elapsed_time:.9f} sec",
        fun_name=module_name, event=dispatch.JAXPR_TO_MLIR_MODULE_EVENT):
    lowering_result = mlir.lower_jaxpr_to_module(
        module_name,
        closed_jaxpr,
        num_const_args=num_const_args,
        ordered_effects=ordered_effects,
        backend=backend,
        platforms=platforms,
        axis_context=axis_ctx,
        in_avals=in_avals,
        out_avals=out_avals,
        donated_args=donated_invars,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove ordered=True (use unordered prints) or drop the debug print entirely in multi-device runs
2. Register the effect as shardable (effects.shardable_ordered_effects) if its semantics permit per-device execution
3. Run the computation on a single device if ordering is essential

Example fix

# before
jax.debug.print('x={}', x, ordered=True)  # inside multi-device pjit

# after
jax.debug.print('x={}', x)  # unordered; or remove entirely

num_devices = len(jax.devices())
if num_devices > 1:
    # ensure no ordered effects in traced function
    jaxpr = jax.make_jaxpr(fn)(*args)
    from jax._src import effects
    bad = effects.ordered_effects.filter_not_in(effects.shardable_ordered_effects).filter_in(jaxpr.jaxpr.effects)
    assert not bad, f'ordered effects unsupported multi-device: {bad}'

When it happens

Trigger: A jitted/pjitted function over a multi-device mesh containing ordered effects — commonly jax.debug.print with ordered=True, or experimental host callback effects — while running with more than one device in the sharding.

Common situations: Debug prints left with ordered=True in multi-device pipelines; scaling single-GPU code to multi-GPU/TPU; custom effect implementations not marked shardable.

Related errors

• Effects not supported in `while`: {}
• to_dlpack can only pack a dlpack tensor from an array on a s
• Query, key and value should have same sharding.
• Effects not supported in `custom_jvp`: {disallowed_effects}
• from_dlpack can only unpack a dlpack tensor onto a singular

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