jax-ml/jax · error · ValueError

Cannot lower jaxpr with effects: {closed_jaxpr.effects}

Error message

Cannot lower jaxpr with effects: {closed_jaxpr.effects}

What it means

When JAX lowers a jaxpr for custom partitioning (pmap-style path), it refuses jaxprs that contain effects which cannot be lowered (anything not in effects_lib.lowerable_effects). Ordered effects like print or state effects have no representation in this lowering path, so JAX rejects the computation up front.

Source

Thrown at jax/_src/interpreters/mlir.py:3339

    dst_symtab.insert(op)

  return renamings["main"]


DEVICE_TO_DEVICE_TYPE = 1
SEND_TO_HOST_TYPE = 2
RECV_FROM_HOST_TYPE = 3

def build_mlir_module_helper(
    closed_jaxpr: core.Jaxpr, *, name: str,
    platforms: Sequence[str],
    backend: xc.Client | None,
    axis_context: AxisContext) -> ir.Module:
  """Helper to generate pmap-style XLA computations for custom partitioners."""
  unlowerable_effects = effects_lib.lowerable_effects.filter_not_in(
      closed_jaxpr.effects)
  if unlowerable_effects:
    raise ValueError(f'Cannot lower jaxpr with effects: {closed_jaxpr.effects}')
  lowering_result = lower_jaxpr_to_module(
      name, closed_jaxpr, num_const_args=0,
      in_avals=closed_jaxpr.in_avals,
      out_avals=closed_jaxpr.out_avals,
      backend=backend, ordered_effects=[],
      donated_args=[False] * len(closed_jaxpr.invars),
      axis_context=axis_context, platforms=platforms,
      lowering_parameters=LoweringParameters(hoist_constants_as_args=False))
  return lowering_result.module


def custom_call(
    call_target_name: str,
    *,
    result_types: Sequence[ir.Type],
    operands: Sequence[ir.Value],
    backend_config: str | bytes | dict[str, ir.Attribute] = "",
    has_side_effect: bool = False,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove or gate effectful operations (e.g. jax.debug.print) from the function being custom-partitioned
2. Move effectful logic outside the partitioned region and pass values in/out instead
3. If you own the primitive, register its effects as lowerable or implement a lowering-compatible path

Example fix

# before
def f(x):
    jax.debug.print('x={}', x)
    return x * 2

# after
def f(x):
    return x * 2

# Before partitioning, inspect for effects:
with jax.make_jaxpr(func) as jaxpr:
    jaxpr_fn = jax.make_jaxpr(func)
jaxpr = jaxpr_fn(*example_args)
assert not jaxpr.jaxpr.effects, f'effects present: {jaxpr.jaxpr.effects}'

try:
    partitioned = spmd.partition(func, ...)
except ValueError as e:
    if 'Cannot lower jaxpr with effects' in str(e):
        # strip debug prints / effects and retry
        ...

When it happens

Trigger: Using custom partitioners (partition.custom_partitioning or pmap-style lowering helpers) on a function whose traced jaxpr contains effects, e.g. jax.debug.print, random state threading implemented as effects, or custom primitives with ordered effects.

Common situations: Leaving debug prints (jax.debug.print) inside a function that is later partitioned with a custom partitioner; adding a new primitive with effects and routing it through the custom-partitioning lowering path.

Related errors

• Cannot lower effectful `reduce_window`.
• multi-platform lowering for buffer_callback
• Nesting `compute_on` with different compute types is not all
• unbound axis name: {axis_name}
• Effects not supported in `custom_jvp`: {disallowed_effects}

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