jax-ml/jax · error · ValueError

Invalid TMEMLayout: {self}

Error message

Invalid TMEMLayout: {self}

What it means

TMEMLayout.to_mgpu() exhausts its match over known TMEMLayout enum members (sparse metadata, M64 collective, scales layouts, etc.) and falls through to this raise. In practice this indicates an enum member added without a corresponding case, i.e. an internal/library bug or an exhaustive-match desync after a partial upgrade.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:1968

  SCALES_LAYOUT = enum.auto()
  SPARSE_METADATA_LAYOUT = enum.auto()
  M64_COLLECTIVE_LAYOUT = enum.auto()
  SCALES_M64_COLLECTIVE_LAYOUT = enum.auto()

  def __call__(self, *args, **kwargs) -> ParameterizedLayout:
    return ParameterizedLayout(self, args, kwargs)

  def to_mgpu(self, *args, **kwargs) -> tcgen05.TMEMLayout:
    match self:
      case TMEMLayout.SCALES_LAYOUT:
        return tcgen05.scales_layout(*args, **kwargs)
      case TMEMLayout.SPARSE_METADATA_LAYOUT:
        return tcgen05.sparse_meta_layout(*args, **kwargs)
      case TMEMLayout.M64_COLLECTIVE_LAYOUT:
        return tcgen05.tmem_m64_collective_layout(*args, **kwargs)
      case TMEMLayout.SCALES_M64_COLLECTIVE_LAYOUT:
        return tcgen05.b_scales_m64_collective_layout(*args, **kwargs)
    raise ValueError(f"Invalid TMEMLayout: {self}")


def TryClusterCancelResult(
    num_buffers: int | None = None) -> pallas_core.MemoryRef:
  """Helper function to create Refs for cluster launch control results.

  Args:
    num_buffers: Optional argument for specifying the number of buffers
      to allocate. If None, will return a single 16-byte buffer. If specified,
      will return a (num_buffers, 16)-shaped buffer.

  Returns:
    A MemoryRef with the correct shape for holding the opaque cluster launch
    control result.
  """
  if num_buffers is None:
    return SMEM((16,), jnp.int8)
  else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reinstall/align jax and jaxlib to the same version (pip install -U jax jaxlib)
2. If you added a TMEMLayout member, add the corresponding match case in to_mgpu()
3. Report upstream if it reproduces on a clean, matched install

Example fix

// before
# mixed versions: jax 0.9.x with jaxlib 0.8.y

// after
pip install -U --force-reinstall jax jaxlib  # same release pair

import jax, jaxlib
# ensure matched versions before using TMEMLayout
assert jax.__version__.split('.')[:2] == jaxlib.__version__.split('.')[:2], 'jax/jaxlib mismatch'

try:
    return tmem_layout.to_mgpu(*args, **kwargs)
except ValueError as e:
    if 'Invalid TMEMLayout' in str(e):
        raise RuntimeError('jax/jaxlib version mismatch or missing TMEMLayout case; reinstall matched versions') from e
    raise

When it happens

Trigger: A new TMEMLayout member exists but to_mgpu() has no case for it; or a stale/partially-updated JAX install mixes an old core.py with new enum definitions.

Common situations: Mixing JAX versions in the same environment (e.g. jax and jaxlib from different releases); running bleeding-edge code where a TMEMLayout was added without updating to_mgpu; monkeypatching the enum.

Related errors

• Unsupported device ID type: {device_id_type}
• Pallas TPU requires a recent libtpu version (at least 0.0.46
• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• Cannot specify both out_shape and out_type. Use out_type.

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