{"record":{"id":"03d8a0572ff9ea3d","repo":"jax-ml/jax","slug":"invalid-tmemlayout-self","errorCode":null,"errorMessage":"Invalid TMEMLayout: {self}","messagePattern":"Invalid TMEMLayout: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":1968,"sourceCode":"  SCALES_LAYOUT = enum.auto()\n  SPARSE_METADATA_LAYOUT = enum.auto()\n  M64_COLLECTIVE_LAYOUT = enum.auto()\n  SCALES_M64_COLLECTIVE_LAYOUT = enum.auto()\n\n  def __call__(self, *args, **kwargs) -> ParameterizedLayout:\n    return ParameterizedLayout(self, args, kwargs)\n\n  def to_mgpu(self, *args, **kwargs) -> tcgen05.TMEMLayout:\n    match self:\n      case TMEMLayout.SCALES_LAYOUT:\n        return tcgen05.scales_layout(*args, **kwargs)\n      case TMEMLayout.SPARSE_METADATA_LAYOUT:\n        return tcgen05.sparse_meta_layout(*args, **kwargs)\n      case TMEMLayout.M64_COLLECTIVE_LAYOUT:\n        return tcgen05.tmem_m64_collective_layout(*args, **kwargs)\n      case TMEMLayout.SCALES_M64_COLLECTIVE_LAYOUT:\n        return tcgen05.b_scales_m64_collective_layout(*args, **kwargs)\n    raise ValueError(f\"Invalid TMEMLayout: {self}\")\n\n\ndef TryClusterCancelResult(\n    num_buffers: int | None = None) -> pallas_core.MemoryRef:\n  \"\"\"Helper function to create Refs for cluster launch control results.\n\n  Args:\n    num_buffers: Optional argument for specifying the number of buffers\n      to allocate. If None, will return a single 16-byte buffer. If specified,\n      will return a (num_buffers, 16)-shaped buffer.\n\n  Returns:\n    A MemoryRef with the correct shape for holding the opaque cluster launch\n    control result.\n  \"\"\"\n  if num_buffers is None:\n    return SMEM((16,), jnp.int8)\n  else:","sourceCodeStart":1950,"sourceCodeEnd":1986,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L1950-L1986","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Reinstall/align jax and jaxlib to the same version (pip install -U jax jaxlib)","If you added a TMEMLayout member, add the corresponding match case in to_mgpu()","Report upstream if it reproduces on a clean, matched install"],"exampleFix":"// before\n# mixed versions: jax 0.9.x with jaxlib 0.8.y\n\n// after\npip install -U --force-reinstall jax jaxlib  # same release pair","handlingStrategy":"retry","validationCode":"import jax, jaxlib\n# ensure matched versions before using TMEMLayout\nassert jax.__version__.split('.')[:2] == jaxlib.__version__.split('.')[:2], 'jax/jaxlib mismatch'","typeGuard":null,"tryCatchPattern":"try:\n    return tmem_layout.to_mgpu(*args, **kwargs)\nexcept ValueError as e:\n    if 'Invalid TMEMLayout' in str(e):\n        raise RuntimeError('jax/jaxlib version mismatch or missing TMEMLayout case; reinstall matched versions') from e\n    raise","preventionTips":["Install jax and jaxlib from the same release","Pin versions in requirements/lockfile","If adding TMEMLayout members, add the to_mgpu case in the same change"],"tags":["jax","pallas","mosaic-gpu","tmem-layout","enum","version-mismatch"],"backgroundTag":"library-version-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}