jax-ml/jax · error · ValueError

Rule cannot handle an op with vector operands or results

Error message

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

What it means

Raised by the func/gpu.launch traversal lowering when an op is expected to carry Mosaic GPU layout/transform attributes but the registered rule cannot handle vector-typed operands or results. Internal layout-lowering limitation for that op.

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

• A address calculation for multiple M tiles
• A scale address calculation for multiple M tiles
• A sparse metadata address calculation for multiple tiles
• address_space not supported
• B scale address calculation for multiple N tiles

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

Appendix: 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)