jax-ml/jax · error · ValueError

layouts passed to `with_layout_constraint` must be of type `

Error message

layouts passed to `with_layout_constraint` must be of type `Layout`. Got {[type(l) for l in layouts_flat]}

What it means

`with_layout_constraint(x, layouts)` requires every leaf of `layouts` to be a `jax.experimental.layout.Layout` instance. Passing raw tuples, strings, or device-layout objects triggers this validation before the layout constraint is applied to the avals.

Source

Thrown at jax/_src/pjit.py:2501

          'Context mesh cannot be empty. Please use `jax.set_mesh` API to enter'
          ' into a mesh context when using `explicit_axes` API.')
    with mesh_lib.use_abstract_mesh(mesh_info.new):
      args = reshard(args, _in_sharding)
      out = fun(*args, **kwargs)
    out_specs = tree_map(lambda o: core.modify_spec_for_auto_manual(
        core.typeof(o).sharding.spec, mesh_lib.get_abstract_mesh()), out)
    return reshard(out, out_specs)
  return decorator

# -------------------- with_layout_constraint --------------------

def with_layout_constraint(x, layouts):
  x_flat, tree = tree_flatten(x)
  x_avals_flat = [core.shaped_abstractify(x) for x in x_flat]
  layouts_flat = tuple(flatten_axes("with_layout_constraint layouts", tree,
                                    layouts))
  if any(not isinstance(l, Layout) for l in layouts_flat):
    raise ValueError(
        'layouts passed to `with_layout_constraint` must be of type'
        f' `Layout`. Got {[type(l) for l in layouts_flat]}')
  check_aval_layout_compatibility(
      layouts_flat, x_avals_flat, ("",) * len(layouts_flat),
      "with_layout_constraint arguments")
  outs = [layout_constraint_p.bind(xf, layout=l)
          for xf, l in zip(x_flat, layouts_flat)]
  return tree_unflatten(tree, outs)

layout_constraint_p = core.Primitive('layout_constraint')
layout_constraint_p.def_abstract_eval(lambda x, **_: x)
ad.deflinear2(layout_constraint_p,
              lambda ct, _, **params: (layout_constraint_p.bind(ct, **params),))

def _layout_constraint_impl(x, *, layout):
  if not isinstance(x, xc.ArrayImpl):
    raise ValueError(
        'with_layout_constraint in eager mode can only be applied to'

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Construct layouts with `jax.experimental.layout.Layout(...)` and pass those instances
2. Verify with `isinstance(l, jax.experimental.layout.Layout)` before calling
3. Update code if upgrading from an older JAX where layouts were tuples

Example fix

# before
with_layout_constraint(x, ((1,0),))
# after
from jax.experimental.layout import Layout
with_layout_constraint(x, Layout((1,0)))

from jax.experimental.layout import Layout
from jax.tree_util import tree_leaves
def all_valid_layouts(layouts) -> bool:
    return all(isinstance(l, Layout) for l in tree_leaves(layouts))

try:
    with_layout_constraint(x, layouts)
except ValueError as e:
    if 'must be of type' in str(e):
        layouts = tree_map(Layout, layouts)  # or fix construction
    else:
        raise

When it happens

Trigger: Calling `jax.experimental.pjit.with_layout_constraint(x, layouts)` with layouts given as e.g. tuples like `(1,0)`, strings, or Layout-like objects from another library or older JAX version.

Common situations: Porting code that used raw layout tuples from XLA; mixing up `Layout` with `sharding` specs; version drift where the accepted layout type changed.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• Expected mode of type `LayoutMode`. Got type: {type(mode)}
• Invalid value received for the layout argument. Expected val
• The layout of ShapedArray should not be `AutoLayout` when la
• Argument '{arg}' of type {type(arg)} is not a valid JAX type
• SymbolicScope constraint must be a string: got {repr(c_str)}

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