{"record":{"id":"dee0f3726080c7ff","repo":"jax-ml/jax","slug":"layouts-passed-to-with-layout-constraint-must-be","errorCode":null,"errorMessage":"layouts passed to `with_layout_constraint` must be of type `Layout`. Got {[type(l) for l in layouts_flat]}","messagePattern":"layouts passed to `with_layout_constraint` must be of type `Layout`\\. Got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pjit.py","lineNumber":2501,"sourceCode":"          'Context mesh cannot be empty. Please use `jax.set_mesh` API to enter'\n          ' into a mesh context when using `explicit_axes` API.')\n    with mesh_lib.use_abstract_mesh(mesh_info.new):\n      args = reshard(args, _in_sharding)\n      out = fun(*args, **kwargs)\n    out_specs = tree_map(lambda o: core.modify_spec_for_auto_manual(\n        core.typeof(o).sharding.spec, mesh_lib.get_abstract_mesh()), out)\n    return reshard(out, out_specs)\n  return decorator\n\n# -------------------- with_layout_constraint --------------------\n\ndef with_layout_constraint(x, layouts):\n  x_flat, tree = tree_flatten(x)\n  x_avals_flat = [core.shaped_abstractify(x) for x in x_flat]\n  layouts_flat = tuple(flatten_axes(\"with_layout_constraint layouts\", tree,\n                                    layouts))\n  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'","sourceCodeStart":2483,"sourceCodeEnd":2519,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pjit.py#L2483-L2519","documentation":"`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.","triggerScenarios":"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.","commonSituations":"Porting code that used raw layout tuples from XLA; mixing up `Layout` with `sharding` specs; version drift where the accepted layout type changed.","solutions":["Construct layouts with `jax.experimental.layout.Layout(...)` and pass those instances","Verify with `isinstance(l, jax.experimental.layout.Layout)` before calling","Update code if upgrading from an older JAX where layouts were tuples"],"exampleFix":"# before\nwith_layout_constraint(x, ((1,0),))\n# after\nfrom jax.experimental.layout import Layout\nwith_layout_constraint(x, Layout((1,0)))","handlingStrategy":"type-guard","validationCode":null,"typeGuard":"from jax.experimental.layout import Layout\nfrom jax.tree_util import tree_leaves\ndef all_valid_layouts(layouts) -> bool:\n    return all(isinstance(l, Layout) for l in tree_leaves(layouts))","tryCatchPattern":"try:\n    with_layout_constraint(x, layouts)\nexcept ValueError as e:\n    if 'must be of type' in str(e):\n        layouts = tree_map(Layout, layouts)  # or fix construction\n    else:\n        raise","preventionTips":["Construct layouts only via Layout(...)","Add an isinstance check helper in test suites"],"tags":["jax","layout","type-validation"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}