jax-ml/jax · error · TypeError

Expected mode of type `LayoutMode`. Got type: {type(mode)}

Error message

Expected mode of type `LayoutMode`. Got type: {type(mode)}

What it means

jax._src.layout.use_layout_mode is a context manager that only accepts jax._src.layout.LayoutMode enum values. Passing a string like 'default' or any other type raises TypeError.

Source

Thrown at jax/_src/layout.py:48

  def __repr__(self):
    return "AutoLayout"
AutoLayout = AutoLayoutSingleton()


class LayoutMode(enum.Enum):
  AUTO = enum.auto()
  JAX = enum.auto()
  PALLAS_TPU = enum.auto()
  PALLAS_GPU = enum.auto()

def get_layout_mode():
  val = jax_config.layout_tracing_mode.value
  return LayoutMode.AUTO if val is None else val

@contextmanager
def use_layout_mode(mode):
  if not isinstance(mode, LayoutMode):
    raise TypeError(
        f'Expected mode of type `LayoutMode`. Got type: {type(mode)}')
  prev_mode = jax_config.layout_tracing_mode.swap_local(mode)
  try:
    yield
  finally:
    jax_config.layout_tracing_mode.set_local(prev_mode)


class Layout:
  major_to_minor: tuple[int, ...]
  tiling: tuple[tuple[int, ...], ...] | None
  sub_byte_element_size_in_bits: int

  AUTO = AutoLayout

  def __init__(self, major_to_minor: tuple[int, ...],
                tiling: tuple[tuple[int, ...], ...] | None = None,
                sub_byte_element_size_in_bits: int = 0):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Import and use LayoutMode members: from jax._src.layout import LayoutMode; use_layout_mode(LayoutMode.AUTO)
2. Map user strings to enum values before calling

Example fix

# before
use_layout_mode('major')
# after
from jax._src.layout import LayoutMode
use_layout_mode(LayoutMode.MAJOR)

from jax._src.layout import LayoutMode
assert isinstance(mode, LayoutMode)

import jax._src.layout as L
def is_layout_mode(m): return isinstance(m, L.LayoutMode)

When it happens

Trigger: Calling use_layout_mode('major') or use_layout_mode(None) instead of a LayoutMode enum member.

Common situations: Custom tracing pipelines or libraries wrapping JAX's layout tracing configuration; passing config strings read from user settings.

Understand the failure class

Background: Invalid enum value errors: "Unknown type", "Invalid scope", "must be one of" — when a string is not on the library's allowed list — this error's family across 23 libraries.

Related errors

• Invalid value received for the layout argument. Expected val
• layouts passed to `with_layout_constraint` must be of type `
• Default value must be of type str, got {default} of type {ge
• Invalid value "{default}" for JAX flag {name}
• new enum value must be in {enum_values}, got {new_val} of ty

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