jax-ml/jax · error · ValueError

Failed to infer a possible set of layouts. This should only

Error message

Failed to infer a possible set of layouts. This should only happen if user-provided layout casts are unsatisfiable.

What it means

The global layout constraint system became Unsatisfiable, meaning no assignment of layouts can satisfy all constraints; per the message this is only expected when user-provided layout casts conflict with each other.

Source

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

              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)

  for op in module.body:
    traverse_op(op, gather_constraints)
    # Short-circuit if we have an unsatisfiable constraint system, we won't
    # construct anything useful anymore.
    if isinstance(global_constraint_system, cs.Unsatisfiable):
      break

  if isinstance(global_constraint_system, cs.Unsatisfiable):
    raise ValueError(
        "Failed to infer a possible set of layouts. This should only happen if "
        "user-provided layout casts are unsatisfiable."
    )

  constraints = derive_relayout_constraints(ctx.value_sites_for_variable)
  global_constraint_system &= cs.ConstraintSystem(constraints=constraints)
  assert not isinstance(global_constraint_system, cs.Unsatisfiable)

  # Add additional (redundant) constraints which helps the search converge
  # faster.
  global_constraint_system = cs.saturate_distinct_from_splat(
      global_constraint_system
  )
  assert not isinstance(global_constraint_system, cs.Unsatisfiable)
  global_constraint_system = (
      cs.canonicalize_strict_non_splat_relayouts_to_equals(
          global_constraint_system
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Remove or relax user-specified layout casts one at a time to find the conflicting pair
Make the casts compatible (chain them via intermediate casts instead of demanding both at once)
Let layout inference choose layouts automatically

Example fix

// before
x = mgpu.layout_cast(x, layout_a)
y = f(mgpu.layout_cast(x, layout_b))
// after
x = mgpu.layout_cast(x, layout_a)
y = f(x)  # single consistent cast

Defensive patterns

Strategy: try-catch

Try / catch

try:
    infer_layouts(module)
except ValueError as e:
    if 'unsatisfiable' in str(e):
        strip_user_casts_and_retry(module)

Prevention

Minimize explicit layout casts; let inference choose
Introduce casts one at a time when tuning layouts

When it happens

Trigger: Two or more layout_casts (or explicit in/out layouts) that impose mutually exclusive constraints on the same value, e.g. requiring a register layout and a shared-memory layout simultaneously.

Common situations: Complex Mosaic kernels where users pin layouts at both ends of a computation that inference cannot reconcile.

Related errors

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