{"record":{"id":"d0757003005a35f0","repo":"jax-ml/jax","slug":"with-layout-constraint-in-eager-mode-can-only-be-a","errorCode":null,"errorMessage":"with_layout_constraint in eager mode can only be applied to jax.Arrays. Got {type(x)}","messagePattern":"with_layout_constraint in eager mode can only be applied to jax\\.Arrays\\. Got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pjit.py","lineNumber":2518,"sourceCode":"  if any(not isinstance(l, Layout) for l in layouts_flat):\n    raise ValueError(\n        'layouts passed to `with_layout_constraint` must be of type'\n        f' `Layout`. Got {[type(l) for l in layouts_flat]}')\n  check_aval_layout_compatibility(\n      layouts_flat, x_avals_flat, (\"\",) * len(layouts_flat),\n      \"with_layout_constraint arguments\")\n  outs = [layout_constraint_p.bind(xf, layout=l)\n          for xf, l in zip(x_flat, layouts_flat)]\n  return tree_unflatten(tree, outs)\n\nlayout_constraint_p = core.Primitive('layout_constraint')\nlayout_constraint_p.def_abstract_eval(lambda x, **_: x)\nad.deflinear2(layout_constraint_p,\n              lambda ct, _, **params: (layout_constraint_p.bind(ct, **params),))\n\ndef _layout_constraint_impl(x, *, layout):\n  if not isinstance(x, xc.ArrayImpl):\n    raise ValueError(\n        'with_layout_constraint in eager mode can only be applied to'\n        f' jax.Arrays. Got {type(x)}')\n  if x.format.layout == layout:\n    return x\n  return api.jit(_identity_fn, out_shardings=Format(layout, x.sharding))(x)\nlayout_constraint_p.def_impl(_layout_constraint_impl)\n\ndef _layout_constraint_hlo_lowering(ctx, x_node, *, layout):\n  aval, = ctx.avals_in\n  out_aval, = ctx.avals_out\n  out = mlir.wrap_with_layout_op(ctx, x_node, out_aval, layout, aval)\n  return [mlir.lower_with_sharding_in_types(ctx, out, out_aval)]\nmlir.register_lowering(layout_constraint_p,\n                       _layout_constraint_hlo_lowering)\n\ndef _layout_constraint_batcher(axis_data, vals_in, dims_in, layout):\n  x, = vals_in\n  d, = dims_in","sourceCodeStart":2500,"sourceCodeEnd":2536,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pjit.py#L2500-L2536","documentation":"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.","triggerScenarios":"Calling `with_layout_constraint` directly (outside `jax.jit`) on a numpy array, a tracer, a donate/grad-traced value, or a Python scalar.","commonSituations":"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.","solutions":["Convert inputs with `jnp.asarray(x)` before calling","Move the `with_layout_constraint` call inside a `jax.jit`-decorated function where abstract evaluation handles tracers","Ensure the operand is a committed jax.Array (e.g. `jax.device_put` it first)"],"exampleFix":"# before\nwith_layout_constraint(np_array, Layout((1,0)))\n# after\nwith_layout_constraint(jnp.asarray(np_array), Layout((1,0)))","handlingStrategy":"type-guard","validationCode":"import jax, jax.numpy as jnp\nx = jnp.asarray(x)  # convert numpy/scalars before eager call","typeGuard":"import jax\ndef is_eager_jax_array(x) -> bool:\n    return isinstance(x, jax.Array) and isinstance(x, jax._src.xla_bridge.ArrayImpl) or type(x).__name__ == 'ArrayImpl'","tryCatchPattern":null,"preventionTips":["Convert with jnp.asarray before eager use","Keep with_layout_constraint inside jit for traced values"],"tags":["jax","layout","eager-mode"],"backgroundTag":"wrong-argument-type-at-runtime","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}