{"record":{"id":"8e890389b1f76b01","repo":"jax-ml/jax","slug":"nesting-compute-on-with-different-compute-types","errorCode":null,"errorMessage":"Nesting `compute_on` with different compute types is not allowed.","messagePattern":"Nesting `compute_on` with different compute types is not allowed\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/compute_on.py","lineNumber":127,"sourceCode":"  return wrapped\n\ncompute_on_p = core.Primitive('compute_on')\ncompute_on_p.multiple_results = True\ndispatch.simple_impl(compute_on_p)\n\n\ndef _compute_on_abstract_eval(*in_avals, jaxpr, compute_type, out_memory_spaces,\n                              compiler_options_json):\n  out_avals = [a.update(memory_space=s) if isinstance(a, core.ShapedArray)\n               else a for a, s in zip(jaxpr.out_avals, out_memory_spaces)]\n  return out_avals, core.positional_effects(jaxpr)\ncompute_on_p.def_effectful_abstract_eval(_compute_on_abstract_eval)\n\n\ndef _compute_on_lowering(ctx, *args, jaxpr, compute_type, out_memory_spaces,\n                         compiler_options_json):\n  if dispatch.jaxpr_has_primitive(jaxpr, 'compute_on'):\n    raise ValueError(\"Nesting `compute_on` with different compute types is \"\n                     \"not allowed.\")\n  const_args_and_avals = core.jaxpr_const_args(jaxpr)\n  const_args, const_avals = unzip2(const_args_and_avals)\n  const_arg_values = [\n      mlir.ir_constants(c, const_lowering=ctx.const_lowering, aval=aval)\n      for c, aval in const_args_and_avals]\n  in_avals = (*const_avals, *ctx.avals_in)\n  func_op, output_types, effects = mlir.lower_called_computation(\n      \"compute_on\", jaxpr, ctx.module_context, len(const_args), in_avals,\n      ctx.avals_out, ctx.tokens_in)\n\n  symbol_name = func_op.name.value\n  flat_output_types, treedef = mlir.ir_tree_registry.flatten(output_types)\n  tokens = [ctx.tokens_in.get(eff) for eff in effects]\n  args = (*ctx.dim_var_values, *tokens, *const_arg_values, *args)\n  flat_args, _ = mlir.ir_tree_registry.flatten(args)\n  call = func_dialect.CallOp(\n      flat_output_types, ir.FlatSymbolRefAttr.get(symbol_name),","sourceCodeStart":109,"sourceCodeEnd":145,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/compute_on.py#L109-L145","documentation":"compute_on is implemented as a JAX primitive whose lowering inspects the jaxpr for a nested compute_on primitive. JAX cannot lower code where one compute_on region with one compute type contains another compute_on with a different compute type, so lowering raises this ValueError. It is a structural restriction on mixing compute types in nested fashion.","triggerScenarios":"Applying @jax.compute_on(compute_type='cpu') to a function that internally calls another function decorated with @jax.compute_on(compute_type='gpu') (or any differing compute_type), then triggering compilation (jit, grad, .lower()).","commonSituations":"Composing library code where an inner helper is already compute_on-decorated with a different backend; refactoring backend placement and accidentally wrapping an already-wrapped function.","solutions":["Remove the inner compute_on decorator and hoist the compute placement so regions don't overlap","Make both nested compute_on calls use the same compute_type","Move the inner compute_on-wrapped call outside the outer wrapped region (call it before/after instead of inside)"],"exampleFix":"# before\n@jax.compute_on(compute_type='gpu')\ndef outer(x):\n    return inner(x)  # inner is @jax.compute_on(compute_type='cpu')\n# after\n@jax.compute_on(compute_type='gpu')\ndef outer(x):\n    return inner_body(x)  # plain function, no inner compute_on","handlingStrategy":"validation","validationCode":"import jax\nimport jax.experimental.compute_on as co\n\n# before composing, check inner functions aren't compute_on-wrapped with a different type\nimport jax._src.compute_on as _co\nassert not getattr(inner_fn, '_compute_on_compute_type', _co_current) != outer_type, 'nested compute_on conflict'","typeGuard":null,"tryCatchPattern":"try:\n    jitted = jax.jit(outer).lower(x).compile()\nexcept ValueError as e:\n    if 'Nesting `compute_on`' in str(e):\n        raise RuntimeError('refactor: hoist inner compute_on out of outer region') from e\n    raise","preventionTips":["Keep one compute_on layer per call region; don't wrap already-wrapped functions","Document which helpers are compute_on-decorated so callers don't re-wrap","Compile eagerly in tests (call .lower().compile()) to catch lowering errors early"],"tags":["jax","compute-on","nesting","lowering"],"backgroundTag":"nested-context-conflict","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}