jax-ml/jax · error · ValueError

A custom return op must terminate the block.

Error message

A custom return op must terminate the block.

What it means

When lowering an mgpu custom primitive, Mosaic relocates the ops from the primitive's body region into the caller and expects the region to end with a custom return op. If none was found, the body is malformed.

Source

Thrown at jax/experimental/mosaic/gpu/dialect_lowering.py:2636

) -> Sequence[ir.Value]:
  """Lowering rule for mgpu.CustomPrimitiveOp."""
  del ctx
  block = op.body.blocks[0]
  for arg, operand in zip(block.arguments, op.operands, strict=True):
    arg.replace_all_uses_with(operand)

  return_op = None
  ip = ir.InsertionPoint.current
  for block_op in block.operations:
    if isinstance(block_op.opview, mgpu.ReturnOp):
      assert return_op is None
      return_op = block_op.opview
      continue
    block_op.detach_from_parent()
    ip.insert(block_op)

  if return_op is None:
    raise ValueError("A custom return op must terminate the block.")

  return return_op.operands


@_register_lowering(mgpu.WarpMapOp)
def _mgpu_warp_map_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.WarpMapOp
) -> Sequence[ir.Value]:
  """Lowering rule for mgpu.WarpMapOp."""
  for a, o in zip(op.body.arguments, op.operands, strict=True):
    a.replace_all_uses_with(o)
  warp_ctx = dataclasses.replace(ctx, thread_semantics=utils.ThreadSubset.WARP)
  # We allow the warps to schedule async copies without synchronizing with
  # other warps, so we need to add a barrier here to make sure all reads and
  # writes have completed.
  if ctx.auto_barriers:
    utils.warpgroup_barrier()
  ip = ir.InsertionPoint.current

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure the custom primitive body block ends with the required return op carrying the outputs
2. Use the provided helper/builder APIs that insert the terminator automatically
3. Validate the region has exactly one terminator before lowering

Example fix

// before
body = build_body_without_terminator(primitive)
// after
body = build_body(primitive)
mgpu.return_(results)  # terminator required at end of body

from jax.experimental.mosaic.gpu import mgpu
terminators = [o for o in body.ops if isinstance(o.opview, mgpu.ReturnOp)]
assert terminators, 'custom primitive body must end with return'

When it happens

Trigger: Defining a custom primitive (warp-level custom op) whose body block does not terminate with mgpu.return / the custom return op — e.g. empty body or body ending in another op.

Common situations: Hand-written custom primitive builders; bugs generating the body IR; forgetting the mandatory terminator when constructing the region programmatically.

Related errors

• {ref} is not a memref.
• {ref} has a memory space {mem_ref_ty.memory_space} that is n
• gpu.launch op not found.
• Mosaic GPU does not yet support AMD ROCm devices. Use ``comp
• {op} has an unsupported layout: {out_layout_attr}

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