{"record":{"id":"aecb45532fd081c9","repo":"jax-ml/jax","slug":"layout-attribute-is-only-defined-for-tmem-refs","errorCode":null,"errorMessage":"layout attribute is only defined for TMEM refs","messagePattern":"layout attribute is only defined for TMEM refs","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":628,"sourceCode":"\n  def _iter(self, tracer):\n    return iter(flatten_ref_union(tracer))\n\n  def _getitem(self, tracer, idx):\n    return list(iter(tracer))[idx]\n\n  def _setitem(self, tracer, idx, value):\n    del tracer, idx, value  # Unused.\n    raise ValueError(\"Ref unions can't be assigned to.\")\n\n  def update(self, inner_aval=None, memory_space=None, kind=None):\n    ref = super().update(inner_aval, memory_space, kind)\n    return AbstractRefUnion(ref.inner_aval, self.refs, self.memory_space)\n\n  @functools.cached_property\n  def layout(self) -> tcgen05.TMEMLayout:\n    if self.memory_space != TMEM:\n      raise ValueError(\"layout attribute is only defined for TMEM refs\")\n    return tcgen05.tmem_default_layout(packing=1)\n\n  @functools.cached_property\n  def collective(self) -> bool:\n    if self.memory_space != TMEM:\n      raise ValueError(\"collective attribute is only defined for TMEM refs\")\n    ref_leaves = jax.tree.leaves(self.refs)\n    first_ref = ref_leaves[0]\n    assert all(ref.collective == first_ref.collective for ref in ref_leaves)\n    return first_ref.collective\n\n  def __eq__(self, other):\n    return (\n        type(self) is type(other)\n        and self.inner_aval == other.inner_aval\n        and self.memory_space == other.memory_space\n        and self.refs == other.refs\n    )","sourceCodeStart":610,"sourceCodeEnd":646,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L610-L646","documentation":"In JAX's Mosaic GPU (Pallas) core, AbstractRefUnion.layout is a cached property that returns the TMEM layout of a union-of-refs aval. Because layout semantics only exist for tensor memory (TMEM) references, accessing it on a union whose memory_space is not TMEM raises ValueError. The layout is computed as tcgen05.tmem_default_layout(packing=1).","triggerScenarios":"Accessing `.layout` on an AbstractRefUnion (or a ref aval produced by RefMap/get_ref_aval) whose memory_space is SMEM or another non-TMEM space, e.g. when inspecting ref avals of SMEM buffers inside a Pallas kernel or during avals-based transforms.","commonSituations":"Writing generic code that introspects ref avals (layout, collective, dtype) for both SMEM and TMEM refs; upgrading JAX versions where AbstractRefUnion gained these attributes; porting pipelines that assumed all refs were TMEM.","solutions":["Check `ref.memory_space == TMEM` before accessing `.layout` (or use `getattr(ref, 'memory_space', None)` dispatch)","Branch on memory_space and only query layout for TMEM refs","If you expected TMEM, verify the ref was actually allocated in TMEM (e.g. via `pl.tmem_ref` / tcgen05 APIs) rather than SMEM"],"exampleFix":"# before\nlayout = ref_aval.layout  # ValueError for SMEM unions\n\n# after\nif ref_aval.memory_space == mosaic_gpu_core.TMEM:\n    layout = ref_aval.layout\nelse:\n    layout = None  # SMEM refs have no layout","handlingStrategy":"validation","validationCode":"if ref_aval.memory_space == mosaic_gpu_core.TMEM:\n    layout = ref_aval.layout\nelse:\n    layout = None","typeGuard":"def is_tmem_ref_aval(a) -> bool:\n    return getattr(a, 'memory_space', None) == mosaic_gpu_core.TMEM","tryCatchPattern":"null","preventionTips":["Branch on memory_space before touching layout/collective","Keep TMEM-only introspection in helpers named for TMEM"],"tags":["jax","pallas","mosaic-gpu","tmem","gpu-kernels"],"backgroundTag":"invalid-attribute-access","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}