jax-ml/jax · error · NotImplementedError

Unsupported op type: {op}

Error message

Unsupported op type: {op}

What it means

During control-flow lowering (for/while bodies) Mosaic flattens layout-tiled values and must rewrite terminator-ish ops (yield, condition). Encountering any other op type inside the rewriter triggers NotImplementedError.

Source

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

      if not isinstance(op, last_op_type):
        # `append` moves the operation.
        new_block.append(op)
        ctx.lower_op(op)
      else:
        assert out_template is None
        layouts = (
            inference_utils.in_layouts(op)
            if inference_utils.has_in_layouts_set(op)
            else []
        )
        if isinstance(op, scf.YieldOp):
          flat_operands, out_template = _flatten_ir_values(op.operands, layouts)
          scf.yield_(flat_operands)
        elif isinstance(op, scf.ConditionOp):
          flat_carry, out_template = _flatten_ir_values(op.args, layouts)
          scf.condition(op.condition, flat_carry)
        else:
          raise NotImplementedError(f"Unsupported op type: {op}")
        op.erase()
  assert out_template is not None
  return out_template


@_register_lowering(scf.ForOp, support_warp_semantics=True)
def _for_op_lowering_rule(
    ctx: LoweringContext, for_op: scf.ForOp
) -> MlirLoweringRuleResult:
  if not inference_utils.should_have_layout(for_op):
    return _traverse_op_lowering_rule(ctx, for_op)
  in_layouts = inference_utils.in_layouts(for_op)
  out_layouts = inference_utils.out_layouts(for_op)
  yield_op = for_op.body.operations[len(for_op.body.operations) - 1]
  yield_layouts = inference_utils.in_layouts(yield_op)
  if in_layouts != out_layouts or in_layouts != yield_layouts:
    raise ValueError("Layout mismatch")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Upgrade (or pin) jax so the Mosaic lowering matches the MLIR dialect version
2. Restructure the kernel to avoid unusual control-flow terminators inside layout-carrying loops
3. Report the op type upstream — it indicates a lowering gap, not user error

try:
    lower(kernel_module)
except NotImplementedError as e:
    if 'Unsupported op type' in str(e):
        # restructure control flow or pin jax version
        ...

When it happens

Trigger: An scf.for/scf.while body containing an op that is neither scf.YieldOp nor scf.ConditionOp at the point the flattening rewriter inspects it — typically newly introduced control-flow terminator variants in newer MLIR/scf dialects.

Common situations: Version skew between jax/mosaic and the bundled MLIR introducing new scf ops; exotic hand-built control flow in a kernel.

Related errors

• Input layouts {in_layouts} do not match yield layouts {yield
• Output layouts {out_layouts} do not match condition layouts
• Rule cannot handle an op with vector operands or results: {o
• Multimem refs are not supported in store_tiled_async
• Replicated dimensions are not supported

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