{"record":{"id":"214859d731ec067f","repo":"jax-ml/jax","slug":"unsupported-op-type-op","errorCode":null,"errorMessage":"Unsupported op type: {op}","messagePattern":"Unsupported op type: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2767,"sourceCode":"      if not isinstance(op, last_op_type):\n        # `append` moves the operation.\n        new_block.append(op)\n        ctx.lower_op(op)\n      else:\n        assert out_template is None\n        layouts = (\n            inference_utils.in_layouts(op)\n            if inference_utils.has_in_layouts_set(op)\n            else []\n        )\n        if isinstance(op, scf.YieldOp):\n          flat_operands, out_template = _flatten_ir_values(op.operands, layouts)\n          scf.yield_(flat_operands)\n        elif isinstance(op, scf.ConditionOp):\n          flat_carry, out_template = _flatten_ir_values(op.args, layouts)\n          scf.condition(op.condition, flat_carry)\n        else:\n          raise NotImplementedError(f\"Unsupported op type: {op}\")\n        op.erase()\n  assert out_template is not None\n  return out_template\n\n\n@_register_lowering(scf.ForOp, support_warp_semantics=True)\ndef _for_op_lowering_rule(\n    ctx: LoweringContext, for_op: scf.ForOp\n) -> MlirLoweringRuleResult:\n  if not inference_utils.should_have_layout(for_op):\n    return _traverse_op_lowering_rule(ctx, for_op)\n  in_layouts = inference_utils.in_layouts(for_op)\n  out_layouts = inference_utils.out_layouts(for_op)\n  yield_op = for_op.body.operations[len(for_op.body.operations) - 1]\n  yield_layouts = inference_utils.in_layouts(yield_op)\n  if in_layouts != out_layouts or in_layouts != yield_layouts:\n    raise ValueError(\"Layout mismatch\")\n","sourceCodeStart":2749,"sourceCodeEnd":2785,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2749-L2785","documentation":"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.","triggerScenarios":"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.","commonSituations":"Version skew between jax/mosaic and the bundled MLIR introducing new scf ops; exotic hand-built control flow in a kernel.","solutions":["Upgrade (or pin) jax so the Mosaic lowering matches the MLIR dialect version","Restructure the kernel to avoid unusual control-flow terminators inside layout-carrying loops","Report the op type upstream — it indicates a lowering gap, not user error"],"exampleFix":null,"handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":"try:\n    lower(kernel_module)\nexcept NotImplementedError as e:\n    if 'Unsupported op type' in str(e):\n        # restructure control flow or pin jax version\n        ...","preventionTips":["Pin jax/xla versions consistent with your kernel code","Avoid exotic scf constructs inside layout-carrying loops"],"tags":["mosaic","gpu","scf","control-flow","not-implemented"],"backgroundTag":"unsupported-op-lowering","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}