jax-ml/jax · error · NotImplementedError

No layout inference rule defined for {op}

Error message

No layout inference rule defined for {op}

What it means

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.

Source

Thrown at jax/experimental/mosaic/gpu/layout_inference.py:2761

  ctx = DerivationContext()

  def gather_constraints(op: Any):
    # Terminator ops are handled directly by the op whose region they belong to.
    # This is because they need to be in sync with their parent op's inputs and
    # outputs---and the parent op's constraints therefore need to take them into
    # account.
    if is_terminator(op):
      return
    should_have_layout = (
        inference_utils.should_have_layout(op)
        or inference_utils.should_have_tmem_layout(op)
        or inference_utils.should_have_transforms(op)
    )
    if not should_have_layout:
      return
    rule = _constraint_system_derivation_rules.get(op.OPERATION_NAME, None)
    if rule is None:
      raise NotImplementedError(f"No layout inference rule defined for {op}")
    rule_result = rule(ctx, op)
    nonlocal global_constraint_system
    constraint_system, mapping = rule_result
    for var, sites in mapping.items():
      assert isinstance(var.key, ValueSite)
      for site in sites:
        if site.memory_space != var.memory_space:
          raise ValueError(
              f"Memory space mismatch between variable and {site}:"
              f" {var.memory_space} != {site.memory_space}."
          )
        if site.shape != var.shape:
          raise ValueError(
              f"Shape mismatch between variable and {site}:"
              f" {var.shape} != {site.shape}."
          )
    global_constraint_system &= constraint_system
    ctx.update(mapping)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Replace the unsupported op with supported Mosaic primitives (tiled loops, mgpu helpers)
2. Upgrade jax — new ops gain inference rules over time
3. For custom ops, register a derivation rule via _constraint_system_derivation_rules or annotate the op to opt out (no transforms/layouts)

When it happens

Trigger: 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.

Common situations: Building custom ops in the mgpu dialect, or using newer MLIR ops with an older jax that lacks the rule.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• {uninitialized_memory} + {str(dtype)}
• Batching over dynamic grid values is not supported yet.
• dma_start not implemented in LoJAX yet.
• DMA partial discharge add=True not yet implemented.
• 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/f799c1f84283e6a4. Report an issue: GitHub.