jax-ml/jax · error · ValueError

with_layout_constraint in eager mode can only be applied to

Error message

with_layout_constraint in eager mode can only be applied to jax.Arrays. Got {type(x)}

What it means

In eager (non-jitted) execution, `with_layout_constraint` can only operate on concrete `jax.Array` objects because it inspects and re-lays-out actual buffers. Applying it to tracers, np.ndarrays, or Python scalars fails this check in the impl rule.

Source

Thrown at jax/_src/pjit.py:2518

  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'
        f' jax.Arrays. Got {type(x)}')
  if x.format.layout == layout:
    return x
  return api.jit(_identity_fn, out_shardings=Format(layout, x.sharding))(x)
layout_constraint_p.def_impl(_layout_constraint_impl)

def _layout_constraint_hlo_lowering(ctx, x_node, *, layout):
  aval, = ctx.avals_in
  out_aval, = ctx.avals_out
  out = mlir.wrap_with_layout_op(ctx, x_node, out_aval, layout, aval)
  return [mlir.lower_with_sharding_in_types(ctx, out, out_aval)]
mlir.register_lowering(layout_constraint_p,
                       _layout_constraint_hlo_lowering)

def _layout_constraint_batcher(axis_data, vals_in, dims_in, layout):
  x, = vals_in
  d, = dims_in

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert inputs with `jnp.asarray(x)` before calling
2. Move the `with_layout_constraint` call inside a `jax.jit`-decorated function where abstract evaluation handles tracers
3. Ensure the operand is a committed jax.Array (e.g. `jax.device_put` it first)

Example fix

# before
with_layout_constraint(np_array, Layout((1,0)))
# after
with_layout_constraint(jnp.asarray(np_array), Layout((1,0)))

import jax, jax.numpy as jnp
x = jnp.asarray(x)  # convert numpy/scalars before eager call

import jax
def is_eager_jax_array(x) -> bool:
    return isinstance(x, jax.Array) and isinstance(x, jax._src.xla_bridge.ArrayImpl) or type(x).__name__ == 'ArrayImpl'

When it happens

Trigger: Calling `with_layout_constraint` directly (outside `jax.jit`) on a numpy array, a tracer, a donate/grad-traced value, or a Python scalar.

Common situations: Debugging layout code outside jit then leaving the eager call in; applying the constraint to a value later passed through `jax.grad`, which turns it into a tracer.

Related errors

• The layout of ShapedArray should not be `AutoLayout` when la
• Expected mode of type `LayoutMode`. Got type: {type(mode)}
• Length of major_to_minor and the rank of the value should ma
• Sharding has to be concrete when layout is of type {type(lay
• Invalid value received for the layout argument. Expected val

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