jax-ml/jax · error · TypeError

Invalid value received for the layout argument. Expected val

Error message

Invalid value received for the layout argument. Expected values are `None`, `Layout.AUTO` or an instance of `Layout`. Got {layout} of type {type(layout)}

What it means

Raised by the constructor of jax._src.layout (Layout device-assignment/validation object) when the `layout` argument is not None, not Layout.AUTO, and not a Layout instance. JAX validates layout eagerly so that invalid layout specs fail at API entry rather than during compilation. The message reports the offending value and its type.

Source

Thrown at jax/_src/layout.py:147

class Format:
  __slots__ = ['layout', 'sharding']

  def __init__(self, layout: LayoutOptions = None,
               sharding: ShardingOptions = None):
    # If layout is concrete and sharding is not, error.
    if isinstance(layout, Layout) and sharding is None:
      raise ValueError(
          'Sharding has to be concrete when layout is of type'
          f' {type(layout)}. Please pass a'
          ' `jax.sharding.NamedSharding` or'
          ' `jax.sharding.SingleDeviceSharding` to the sharding argument. Got'
          f' sharding {sharding}'
      )
    if not isinstance(
        layout, (Layout, type(None), AutoLayoutSingleton)):
      raise TypeError(
          'Invalid value received for the layout argument.'
          ' Expected values are `None`, `Layout.AUTO` or an'
          f' instance of `Layout`. Got {layout} of'
          f' type {type(layout)}'
      )
    if not isinstance(sharding, (Sharding, type(None))):
      raise TypeError(
          'Invalid value received for the sharding argument. Expected values'
          ' are `None` or an instance of `jax.Sharding`. Got'
          f' {sharding} of type {type(sharding)}')

    self.layout = layout
    self.sharding = sharding

  def __repr__(self):
    return f'Format(layout={self.layout}, sharding={self.sharding})'

  def __hash__(self):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass jax.sharding.Layout.AUTO (or None) or a real jax.sharding.Layout instance
2. Check you are not passing a PartitionSpec/Sharding into the layout slot (swap arguments)
3. Upgrade/downgrade JAX to a version whose Layout API matches your code
4. If subclassing Layout, ensure you inherit from jax.sharding.Layout, not a duck-typed imitation

Example fix

# before
layout = 'auto'
jax.jit(f, layout=layout)

# after
from jax.sharding import Layout
layout = Layout.AUTO  # or None, or Layout(...)

from jax.sharding import Layout
ok = layout is None or layout is Layout.AUTO or isinstance(layout, Layout)

def is_valid_layout(x) -> bool:
    from jax.sharding import Layout
    return x is None or x is Layout.AUTO or isinstance(x, Layout)

try:
    op(layout=layout)
except TypeError as e:
    if 'layout argument' in str(e):
        layout = None; op(layout=layout)
    else: raise

When it happens

Trigger: Calling an API that takes a layout kwarg (e.g. jax.jit sharding/layout or making a Layout object) with a string like 'auto', a PartitionSpec, a Sharding, or any custom object instead of Layout.AUTO / None / a Layout instance.

Common situations: Porting code that passed layout-like strings; confusing sharding and layout arguments; building a custom Layout subclass hierarchy where the isinstance check no longer matches; older JAX versions where the argument did not exist or accepted anything.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• Expected mode of type `LayoutMode`. Got type: {type(mode)}
• layouts passed to `with_layout_constraint` must be of type `
• The layout of ShapedArray should not be `AutoLayout` when la
• Argument '{arg}' of type {type(arg)} is not a valid JAX type
• SymbolicScope constraint must be a string: got {repr(c_str)}

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