jax-ml/jax · error · TypeError

Missing required keyword argument: 'in_layouts'

Error message

Missing required keyword argument: 'in_layouts'

What it means

The `explicit_layout` decorator requires an `in_layouts` specification; if it wasn't given at decoration time it must be passed as a keyword at call time. Without it JAX cannot relayout the arguments before invoking the function.

Source

Thrown at jax/_src/pjit.py:2639

  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)
  return _explicit_layout(f, **kwargs)

def _explicit_layout(fun, *, in_layouts):
  def decorator(*args, **kwargs):
    if in_layouts is None:
      if "in_layouts" in kwargs:
        _in_layouts = kwargs.pop("in_layouts")
      else:
        raise TypeError("Missing required keyword argument: 'in_layouts'")
    else:
      _in_layouts = in_layouts
    args = relayout(args, _in_layouts)
    mode = get_layout_mode_from_args(args)
    with use_layout_mode(mode):
      out = fun(*args)
    return relayout(out, AutoLayout)
  return decorator


def relayout(xs, out_layouts):
  x_flat, treedef = tree_flatten(xs)
  layouts_flat = flatten_axis_resources(
      "relayout out_layouts", treedef, out_layouts, tupled_args=True)
  out_flat = [relayout_p.bind(x, dst_layout=l)
              for x, l in zip(x_flat, layouts_flat)]
  return tree_unflatten(treedef, out_flat)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass `in_layouts=...` to the decorator
2. Or pass `in_layouts` as a keyword argument when calling the decorated function
3. Match the pytree structure of in_layouts to the function arguments

Example fix

# before
f = explicit_layout(fun)
f(x)
# after
f = explicit_layout(fun, in_layouts=Layout((1,0)))
f(x)

assert in_layouts is not None or 'in_layouts' in call_kwargs, 'in_layouts required'

try:
    f(x)
except TypeError as e:
    if 'in_layouts' in str(e):
        f(x, in_layouts=ly)
    else:
        raise

When it happens

Trigger: Using `explicit_layout(fun)` or calling the decorated function without `in_layouts=...` in either place.

Common situations: Omitting the parameter when experimenting with the new layout API; refactoring away the layouts argument and forgetting call sites that pass it via kwargs.

Understand the failure class

Background: "missing required argument" and "the following required arguments were not provided": what required-argument errors mean and how to fix them — this error's family across 20 libraries.

Related errors

• Context manager for {state.__name__} config option requires
• 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

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