{"record":{"id":"1955c8ae0530c38f","repo":"jax-ml/jax","slug":"cannot-lower-jaxpr-with-effects-closed-jaxpr-eff","errorCode":null,"errorMessage":"Cannot lower jaxpr with effects: {closed_jaxpr.effects}","messagePattern":"Cannot lower jaxpr with effects: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/interpreters/mlir.py","lineNumber":3339,"sourceCode":"    dst_symtab.insert(op)\n\n  return renamings[\"main\"]\n\n\nDEVICE_TO_DEVICE_TYPE = 1\nSEND_TO_HOST_TYPE = 2\nRECV_FROM_HOST_TYPE = 3\n\ndef build_mlir_module_helper(\n    closed_jaxpr: core.Jaxpr, *, name: str,\n    platforms: Sequence[str],\n    backend: xc.Client | None,\n    axis_context: AxisContext) -> ir.Module:\n  \"\"\"Helper to generate pmap-style XLA computations for custom partitioners.\"\"\"\n  unlowerable_effects = effects_lib.lowerable_effects.filter_not_in(\n      closed_jaxpr.effects)\n  if unlowerable_effects:\n    raise ValueError(f'Cannot lower jaxpr with effects: {closed_jaxpr.effects}')\n  lowering_result = lower_jaxpr_to_module(\n      name, closed_jaxpr, num_const_args=0,\n      in_avals=closed_jaxpr.in_avals,\n      out_avals=closed_jaxpr.out_avals,\n      backend=backend, ordered_effects=[],\n      donated_args=[False] * len(closed_jaxpr.invars),\n      axis_context=axis_context, platforms=platforms,\n      lowering_parameters=LoweringParameters(hoist_constants_as_args=False))\n  return lowering_result.module\n\n\ndef custom_call(\n    call_target_name: str,\n    *,\n    result_types: Sequence[ir.Type],\n    operands: Sequence[ir.Value],\n    backend_config: str | bytes | dict[str, ir.Attribute] = \"\",\n    has_side_effect: bool = False,","sourceCodeStart":3321,"sourceCodeEnd":3357,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/interpreters/mlir.py#L3321-L3357","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Remove or gate effectful operations (e.g. jax.debug.print) from the function being custom-partitioned","Move effectful logic outside the partitioned region and pass values in/out instead","If you own the primitive, register its effects as lowerable or implement a lowering-compatible path"],"exampleFix":"# before\ndef f(x):\n    jax.debug.print('x={}', x)\n    return x * 2\n\n# after\ndef f(x):\n    return x * 2","handlingStrategy":"validation","validationCode":"# Before partitioning, inspect for effects:\nwith jax.make_jaxpr(func) as jaxpr:\n    jaxpr_fn = jax.make_jaxpr(func)\njaxpr = jaxpr_fn(*example_args)\nassert not jaxpr.jaxpr.effects, f'effects present: {jaxpr.jaxpr.effects}'","typeGuard":null,"tryCatchPattern":"try:\n    partitioned = spmd.partition(func, ...)\nexcept ValueError as e:\n    if 'Cannot lower jaxpr with effects' in str(e):\n        # strip debug prints / effects and retry\n        ...","preventionTips":["Keep jax.debug.print out of functions destined for custom partitioners","Assert jaxpr effects are empty in tests for partitioned code"],"tags":["jax","effects","custom-partitioning","lowering","pmap"],"backgroundTag":"unsupported-operation-for-codepath","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}