{"record":{"id":"f799c1f84283e6a4","repo":"jax-ml/jax","slug":"no-layout-inference-rule-defined-for-op","errorCode":null,"errorMessage":"No layout inference rule defined for {op}","messagePattern":"No layout inference rule defined for (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/layout_inference.py","lineNumber":2761,"sourceCode":"  ctx = DerivationContext()\n\n  def gather_constraints(op: Any):\n    # Terminator ops are handled directly by the op whose region they belong to.\n    # This is because they need to be in sync with their parent op's inputs and\n    # outputs---and the parent op's constraints therefore need to take them into\n    # account.\n    if is_terminator(op):\n      return\n    should_have_layout = (\n        inference_utils.should_have_layout(op)\n        or inference_utils.should_have_tmem_layout(op)\n        or inference_utils.should_have_transforms(op)\n    )\n    if not should_have_layout:\n      return\n    rule = _constraint_system_derivation_rules.get(op.OPERATION_NAME, None)\n    if rule is None:\n      raise NotImplementedError(f\"No layout inference rule defined for {op}\")\n    rule_result = rule(ctx, op)\n    nonlocal global_constraint_system\n    constraint_system, mapping = rule_result\n    for var, sites in mapping.items():\n      assert isinstance(var.key, ValueSite)\n      for site in sites:\n        if site.memory_space != var.memory_space:\n          raise ValueError(\n              f\"Memory space mismatch between variable and {site}:\"\n              f\" {var.memory_space} != {site.memory_space}.\"\n          )\n        if site.shape != var.shape:\n          raise ValueError(\n              f\"Shape mismatch between variable and {site}:\"\n              f\" {var.shape} != {site.shape}.\"\n          )\n    global_constraint_system &= constraint_system\n    ctx.update(mapping)","sourceCodeStart":2743,"sourceCodeEnd":2779,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/layout_inference.py#L2743-L2779","documentation":"Mosaic's layout inference dispatches on op name via a registry of derivation rules (_constraint_system_derivation_rules). Encountering an op that should have layouts but has no registered rule raises NotImplementedError — the op simply isn't supported by layout inference yet.","triggerScenarios":"Using an mgpu/MLIR op (with vector operands/results or transforms) inside a Mosaic kernel for which no layout inference rule exists, e.g. a newly added or exotic op.","commonSituations":"Building custom ops in the mgpu dialect, or using newer MLIR ops with an older jax that lacks the rule.","solutions":["Replace the unsupported op with supported Mosaic primitives (tiled loops, mgpu helpers)","Upgrade jax — new ops gain inference rules over time","For custom ops, register a derivation rule via _constraint_system_derivation_rules or annotate the op to opt out (no transforms/layouts)"],"exampleFix":null,"handlingStrategy":null,"validationCode":null,"typeGuard":null,"tryCatchPattern":null,"preventionTips":[],"tags":["jax","mosaic","not-implemented","unsupported-op"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}