jax-ml/jax · error · ValueError

Rule cannot handle an op with vector operands or results: {o

Error message

Rule cannot handle an op with vector operands or results: {op}

What it means

The generic traversal lowering rule (used for func.func, gpu.launch, etc.) cannot handle ops whose operands or results carry vector layouts — it would silently drop layout information, so it refuses.

Source

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

  utils.try_cluster_cancel(op.cancellation_result, barrier.barrier_ref, predicate)
  return []


@_register_lowering(mgpu.QueryClusterCancelOp)
def _query_cluster_cancel_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.QueryClusterCancelOp
) -> Sequence[ir.Value]:
  del ctx
  return utils.query_cluster_cancel(op.cancellation_result)


@_register_lowering(func.FuncOp)
@_register_lowering(gpu.LaunchOp)
def _traverse_op_lowering_rule(
    ctx: LoweringContext, op: ir.OpView
) -> MlirLoweringRuleResult:
  if inference_utils.should_have_layout(op):
    raise ValueError(
        f"Rule cannot handle an op with vector operands or results: {op}"
    )
  for region in op.operation.regions:
    for block in region:
      for block_op in list(block):
        with ir.InsertionPoint(block_op):
          ctx.lower_op(block_op)
  return RECURSED


def _should_lower(op: ir.OpView) -> bool:
  """Returns 'true' if the operation should be lowered."""
  return (
      # pyrefly: ignore[missing-attribute]
      op.OPERATION_NAME.startswith("mosaic_gpu.")
      or inference_utils.should_have_layout(op)
      or inference_utils.should_have_transforms(op)
      or inference_utils.should_have_tmem_layout(op)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Register a dedicated lowering rule for the op via _register_lowering (or use an existing op with a rule)
2. Wrap the vector computation in ops that Mosaic knows how to lower
3. Check the op's operands/results — if it's purely scalar it shouldn't have layouts; fix layout inference attachment

if inference_utils.should_have_layout(op):
    raise ValueError(f'{op} needs a dedicated lowering rule; not supported by traversal')

def needs_special_lowering(op) -> bool:
    from jax.experimental.mosaic.gpu import inference_utils
    return inference_utils.should_have_layout(op)

When it happens

Trigger: A layout-annotated vector op reaching _traverse_op_lowering_rule because no specialized lowering was registered for it — e.g. a new or custom op with tiled tensor operands that lacks a registered lowering rule.

Common situations: Custom ops or new dialect ops added without a registered Mosaic lowering; version mismatch where a registered rule is missing.

Related errors

• {ref} has a layout {layout_attr} that does not match the exp
• Unsupported op type: {op}
• Input layouts {in_layouts} do not match yield layouts {yield
• Output layouts {out_layouts} do not match condition layouts
• Multimem refs are not supported in store_tiled_async

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