{"record":{"id":"afe7b4b566478c3b","repo":"jax-ml/jax","slug":"all-aliased-refs-must-have-the-same-memory-space","errorCode":null,"errorMessage":"All aliased Refs must have the same memory space (SMEM or TMEM). Got {(ref.memory_space for ref in ref_leaves)}.","messagePattern":"All aliased Refs must have the same memory space \\(SMEM or TMEM\\)\\. Got (.+?)\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":704,"sourceCode":"      object.__setattr__(self, \"refs\", refs)\n      max_cols = max(map(_ref_group_tmem_col_size, self.refs))\n      is_collective = ref_leaves[0].collective\n      if any(r.collective != is_collective for r in ref_leaves):\n        raise ValueError(\n            \"Some aliased TMEM references are collective and some are not.\"\n        )\n      super().__init__(\n          inner_aval=jax_core.ShapedArray(\n              shape=(128, max_cols,),\n              dtype=jnp.int32,\n          ),\n          memory_space=TMEM,\n          transforms=(),\n          layout=tcgen05.tmem_default_layout(packing=1),\n          collective=all(ref.collective for ref in ref_leaves),\n      )\n    else:\n      raise NotImplementedError(\n          \"All aliased Refs must have the same memory space (SMEM or TMEM). \"\n          f\"Got {(ref.memory_space for ref in ref_leaves)}.\")\n\n  def get_ref_aval(self) -> AbstractRefUnion:\n    inner_aval = jax.core.ShapedArray(self.shape, self.dtype)\n    refs_aval = jax.tree.map(lambda ref: ref.get_ref_aval(), self.refs)\n    return AbstractRefUnion(inner_aval, refs_aval,\n                            memory_space=self.memory_space)\n\n\nIndex = mgpu.DynamicSlice | slice | int | ir.Value\n\n\n@dataclasses.dataclass(frozen=True)\nclass TilingTransform(state_types.Transform):\n  \"\"\"Represents a tiling transformation for memory refs.\n\n  A tiling of (X, Y) on an array of shape (M, N) will result in a transformed","sourceCodeStart":686,"sourceCodeEnd":722,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L686-L722","documentation":"RefMap in Mosaic GPU only supports aliasing refs that all live in the same memory space — either all SMEM or all TMEM. Mixing memory spaces (or including refs in other spaces) in one aliased union raises NotImplementedError. Note the message formats a generator expression, so the printed 'Got' value shows the generator repr, not the actual spaces.","triggerScenarios":"Constructing a RefMap whose leaves contain e.g. one SMEM ref and one TMEM ref, or a ref in an unsupported space (e.g. HBM), so neither the all-SMEM nor the all-TMEM branch of __init__ applies.","commonSituations":"Aliasing scratch buffers across memory spaces in Pallas kernels (e.g. sharing an SMEM staging buffer with a TMEM accumulator); refactors of buffer reuse/aliasing passes; passing a mixed pytree where a RefMap accidentally includes an extra ref.","solutions":["Split the aliased refs into separate RefMaps per memory space","If cross-space sharing was intended, remove it — copy explicitly between SMEM and TMEM instead of aliasing","Debug by printing `[ref.memory_space for ref in jax.tree.leaves(refs)]` before constructing the RefMap (the error message itself is unhelpful due to the generator formatting bug)"],"exampleFix":"# before\nref_map = RefMap((smem_ref, tmem_ref))  # NotImplementedError\n\n# after\nsmem_map = RefMap((smem_ref, other_smem_ref))\ntmem_map = RefMap((tmem_ref, other_tmem_ref))","handlingStrategy":"validation","validationCode":"spaces = {r.memory_space for r in jax.tree.leaves(refs)}\nassert len(spaces) == 1, f'mixed memory spaces: {spaces}'","typeGuard":"def refs_share_memory_space(refs) -> bool:\n    spaces = {r.memory_space for r in jax.tree.leaves(refs)}\n    return len(spaces) == 1","tryCatchPattern":"null","preventionTips":["Alias only within one memory space","Print [r.memory_space for r in leaves] when debugging — the error message hides them"],"tags":["jax","pallas","mosaic-gpu","smem","tmem","aliasing"],"backgroundTag":"unsupported-operation-mixed-types","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}