jax-ml/jax · error · TypeError

All args passed to `explicit_layout` must have the same type

Error message

All args passed to `explicit_layout` must have the same type of layout. Got {layouts=}

What it means

`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.

Source

Thrown at jax/_src/pjit.py:2612

      cur_inps = map(read, eqn.invars)
      if prev_eqn is not None:
        prev_outs = map(read, prev_eqn.outvars)
        # TODO(yashkatariya): Maybe dedup prev_outs and cur_inps.
        prev_outs, cur_inps = optimization_barrier((prev_outs, cur_inps))
        eqn_write(prev_eqn, prev_outs)
      ans = eqn.primitive.bind(*cur_inps, **bind_params)
    eqn_write(eqn, ans)
    prev_eqn = eqn
    core.clean_up_dead_vars(eqn, env, last_used)
  outvals = map(read, jaxpr.outvars)
  return outvals

# ----------------------------- explicit layout --------------------------------

def get_layout_mode_from_args(args):
  layouts = [core.typeof(a).layout for a in args]
  if not all(type(l) is type(layouts[0]) for l in layouts):
    raise TypeError(
        'All args passed to `explicit_layout` must have the same type of'
        f' layout. Got {layouts=}')
  l = layouts[0]
  if isinstance(l, Layout):
    return LayoutMode.JAX
  # TODO(yashkatariya): Replace this with `isinstance(l, ArrayLayout)`.
  elif type(l).__name__ == 'ArrayLayout':
    return LayoutMode.PALLAS_TPU
  elif type(l).__name__ == 'GPUTiledLayout':
    return LayoutMode.PALLAS_GPU
  else:
    return LayoutMode.AUTO


def explicit_layout(f=None, /, *, in_layouts=None):
  kwargs = dict(in_layouts=in_layouts)
  if f is None:
    return lambda g: _explicit_layout(g, **kwargs)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Apply `relayout` (or the same layout mechanism) to ALL arguments so they share one layout type
2. Convert numpy inputs to jax arrays with the same layout before the call
3. Check `jax.core.typeof(a).layout` for each arg to find the odd one out

Example fix

# before
f(numpy_x, relayouted_y)
# after
f(relayout(numpy_x, ly), relayouted_y)

import jax
layouts = [jax.core.typeof(a).layout for a in args]
assert len({type(l) for l in layouts}) == 1, f'mixed layout types: {[type(l) for l in layouts]}'

When it happens

Trigger: 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).

Common situations: 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.

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/781b6ae32dfa6c9c. Report an issue: GitHub.