jax-ml/jax · error · ValueError

Cannot lower jaxpr with 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.

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`.
• accessing .backend in multi-lowering setting. This can…
• axis_index_groups must all be the same size for TPU lowering
• Cannot lower a pallas_call with constants.
• Convolutions with non-static strides, dilation…

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

Appendix: 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)