{"record":{"id":"781b6ae32dfa6c9c","repo":"jax-ml/jax","slug":"all-args-passed-to-explicit-layout-must-have-the","errorCode":null,"errorMessage":"All args passed to `explicit_layout` must have the same type of layout. Got {layouts=}","messagePattern":"All args passed to `explicit_layout` must have the same type of layout\\. Got (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/pjit.py","lineNumber":2612,"sourceCode":"      cur_inps = map(read, eqn.invars)\n      if prev_eqn is not None:\n        prev_outs = map(read, prev_eqn.outvars)\n        # TODO(yashkatariya): Maybe dedup prev_outs and cur_inps.\n        prev_outs, cur_inps = optimization_barrier((prev_outs, cur_inps))\n        eqn_write(prev_eqn, prev_outs)\n      ans = eqn.primitive.bind(*cur_inps, **bind_params)\n    eqn_write(eqn, ans)\n    prev_eqn = eqn\n    core.clean_up_dead_vars(eqn, env, last_used)\n  outvals = map(read, jaxpr.outvars)\n  return outvals\n\n# ----------------------------- explicit layout --------------------------------\n\ndef get_layout_mode_from_args(args):\n  layouts = [core.typeof(a).layout for a in args]\n  if not all(type(l) is type(layouts[0]) for l in layouts):\n    raise TypeError(\n        'All args passed to `explicit_layout` must have the same type of'\n        f' layout. Got {layouts=}')\n  l = layouts[0]\n  if isinstance(l, Layout):\n    return LayoutMode.JAX\n  # TODO(yashkatariya): Replace this with `isinstance(l, ArrayLayout)`.\n  elif type(l).__name__ == 'ArrayLayout':\n    return LayoutMode.PALLAS_TPU\n  elif type(l).__name__ == 'GPUTiledLayout':\n    return LayoutMode.PALLAS_GPU\n  else:\n    return LayoutMode.AUTO\n\n\ndef explicit_layout(f=None, /, *, in_layouts=None):\n  kwargs = dict(in_layouts=in_layouts)\n  if f is None:\n    return lambda g: _explicit_layout(g, **kwargs)","sourceCodeStart":2594,"sourceCodeEnd":2630,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pjit.py#L2594-L2630","documentation":"`explicit_layout` inspects the runtime layout type of all arguments to decide which layout mode to use; mixing args whose layouts are different types (e.g. one with a `Layout` and another with an auto/numpy layout) is ambiguous and rejected.","triggerScenarios":"Calling a function decorated with explicit_layout where one argument carries a JAX `Layout` and another has a default/auto layout or `ArrayLayout` type (e.g. mixing `relayout`-ed arrays with plain numpy-backed arrays).","commonSituations":"Gradually adopting explicit layouts in a codebase where some inputs are converted and others are not; passing a mix of device arrays and numpy arrays.","solutions":["Apply `relayout` (or the same layout mechanism) to ALL arguments so they share one layout type","Convert numpy inputs to jax arrays with the same layout before the call","Check `jax.core.typeof(a).layout` for each arg to find the odd one out"],"exampleFix":"# before\nf(numpy_x, relayouted_y)\n# after\nf(relayout(numpy_x, ly), relayouted_y)","handlingStrategy":"validation","validationCode":"import jax\nlayouts = [jax.core.typeof(a).layout for a in args]\nassert len({type(l) for l in layouts}) == 1, f'mixed layout types: {[type(l) for l in layouts]}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Relayout all inputs uniformly before explicit_layout calls","Log typeof(a).layout per arg during development"],"tags":["jax","layout","consistency"],"backgroundTag":"inconsistent-argument-types","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}