{"record":{"id":"f28ab5fdbbb6065e","repo":"jax-ml/jax","slug":"can-t-transpose-a-tmem-reference","errorCode":null,"errorMessage":"Can't transpose a TMEM reference.","messagePattern":"Can't transpose a TMEM reference\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":1129,"sourceCode":"    collective_axes = (collective_axes,)\n  if not isinstance(ref, pallas_core.TransformedRef):\n    if not isinstance(jax_core.typeof(ref), state_types.AbstractRef):\n      raise TypeError(\"ref must be a reference\")\n    ref = pallas_core.TransformedRef(ref, transforms=())\n  if any(isinstance(t, PeerMemRef) for t in ref.transforms):\n    raise ValueError(\"Can't make a peer reference into a multicast reference.\")\n  return pallas_core.TransformedRef(\n      ref.ref, (*ref.transforms, MulticastRef(collective_axes)),\n  )\n\n\ndef transpose_ref(\n    ref: pallas_core.TransformedRef | Any,\n    permutation: tuple[int, ...],\n) -> pallas_core.TransformedRef:\n  assert hasattr(ref, \"memory_space\")\n  if ref.memory_space == MemorySpace.TMEM:\n    raise ValueError(\"Can't transpose a TMEM reference.\")\n  return ref.transpose(permutation)\n\n\n@tree_util.register_dataclass\n@dataclasses.dataclass(frozen=True)\nclass ExtractAliasedRef(state_types.Transform):\n  \"\"\"Bitcasts the underlying ref at the given offset to the given shape and dtype.\"\"\"\n  dtype: dtypes.DType = jax.tree.static()\n  shape: tuple[int, ...] = jax.tree.static()\n  offset: int = jax.tree.static()\n\n  # The index of the group of this aliased ref within the input RefUnion.\n  alias_group_idx: int = jax.tree.static()\n\n  # TMEM-specific params\n  layout: tcgen05.TMEMLayout | None = jax.tree.static()\n\n  @classmethod","sourceCodeStart":1111,"sourceCodeEnd":1147,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L1111-L1147","documentation":"transpose_ref rejects transposing references whose memory_space is TMEM (tensor memory). TMEM layout on TPU cannot be permuted as an address transform, so the operation raises ValueError.","triggerScenarios":"Calling transpose_ref on a ref allocated in MemorySpace.TMEM (e.g. a tmem_t allocation or TMEM block) with any permutation.","commonSituations":"Kernels using tensor-core TMEM buffers (tcgen05/tmem) where generic transpose helper code is applied to all refs indiscriminately.","solutions":["Skip transposition for TMEM refs: check ref.memory_space first","Use explicit TMEM-aware layout ops (e.g. load/store with transposed indices) instead of transpose_ref","Move data out of TMEM into SMEM/VMEM before transposing"],"exampleFix":"# before\nref = transpose_ref(tmem_ref, (1, 0))\n\n# after\nif ref.memory_space != MemorySpace.TMEM:\n    ref = transpose_ref(ref, (1, 0))","handlingStrategy":"validation","validationCode":"from jax._src.pallas.mosaic_gpu.core import MemorySpace\nif ref.memory_space == MemorySpace.TMEM: raise ValueError('cannot transpose TMEM')","typeGuard":"def is_transposable(ref): return getattr(ref, 'memory_space', None) != MemorySpace.TMEM","tryCatchPattern":null,"preventionTips":["Gate transpose helpers on memory_space checks"],"tags":["jax","pallas","mosaic-gpu","tmem","transpose"],"backgroundTag":"unsupported-memory-space-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}