jax-ml/jax · error · ValueError

Nesting `compute_on` with different compute types is not all

Error message

Nesting `compute_on` with different compute types is not allowed.

What it means

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.

Source

Thrown at jax/_src/compute_on.py:127

  return wrapped

compute_on_p = core.Primitive('compute_on')
compute_on_p.multiple_results = True
dispatch.simple_impl(compute_on_p)


def _compute_on_abstract_eval(*in_avals, jaxpr, compute_type, out_memory_spaces,
                              compiler_options_json):
  out_avals = [a.update(memory_space=s) if isinstance(a, core.ShapedArray)
               else a for a, s in zip(jaxpr.out_avals, out_memory_spaces)]
  return out_avals, core.positional_effects(jaxpr)
compute_on_p.def_effectful_abstract_eval(_compute_on_abstract_eval)


def _compute_on_lowering(ctx, *args, jaxpr, compute_type, out_memory_spaces,
                         compiler_options_json):
  if dispatch.jaxpr_has_primitive(jaxpr, 'compute_on'):
    raise ValueError("Nesting `compute_on` with different compute types is "
                     "not allowed.")
  const_args_and_avals = core.jaxpr_const_args(jaxpr)
  const_args, const_avals = unzip2(const_args_and_avals)
  const_arg_values = [
      mlir.ir_constants(c, const_lowering=ctx.const_lowering, aval=aval)
      for c, aval in const_args_and_avals]
  in_avals = (*const_avals, *ctx.avals_in)
  func_op, output_types, effects = mlir.lower_called_computation(
      "compute_on", jaxpr, ctx.module_context, len(const_args), in_avals,
      ctx.avals_out, ctx.tokens_in)

  symbol_name = func_op.name.value
  flat_output_types, treedef = mlir.ir_tree_registry.flatten(output_types)
  tokens = [ctx.tokens_in.get(eff) for eff in effects]
  args = (*ctx.dim_var_values, *tokens, *const_arg_values, *args)
  flat_args, _ = mlir.ir_tree_registry.flatten(args)
  call = func_dialect.CallOp(
      flat_output_types, ir.FlatSymbolRefAttr.get(symbol_name),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the inner compute_on decorator and hoist the compute placement so regions don't overlap
2. Make both nested compute_on calls use the same compute_type
3. Move the inner compute_on-wrapped call outside the outer wrapped region (call it before/after instead of inside)

Example fix

# before
@jax.compute_on(compute_type='gpu')
def outer(x):
    return inner(x)  # inner is @jax.compute_on(compute_type='cpu')
# after
@jax.compute_on(compute_type='gpu')
def outer(x):
    return inner_body(x)  # plain function, no inner compute_on

import jax
import jax.experimental.compute_on as co

# before composing, check inner functions aren't compute_on-wrapped with a different type
import jax._src.compute_on as _co
assert not getattr(inner_fn, '_compute_on_compute_type', _co_current) != outer_type, 'nested compute_on conflict'

try:
    jitted = jax.jit(outer).lower(x).compile()
except ValueError as e:
    if 'Nesting `compute_on`' in str(e):
        raise RuntimeError('refactor: hoist inner compute_on out of outer region') from e
    raise

When it happens

Trigger: 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()).

Common situations: 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.

Related errors

• multi-platform lowering for buffer_callback
• `compute_on`'s compute_type argument must be a string.
• accessing .backend in multi-lowering setting. This can occur
• the platform for the specified backend {xb.canonicalize_plat
• Cannot lower jaxpr with effects: {closed_jaxpr.effects}

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