jax-ml/jax · error · ValueError

layout attribute is only defined for TMEM refs

Error message

layout attribute is only defined for TMEM refs

What it means

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).

Source

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

  def _iter(self, tracer):
    return iter(flatten_ref_union(tracer))

  def _getitem(self, tracer, idx):
    return list(iter(tracer))[idx]

  def _setitem(self, tracer, idx, value):
    del tracer, idx, value  # Unused.
    raise ValueError("Ref unions can't be assigned to.")

  def update(self, inner_aval=None, memory_space=None, kind=None):
    ref = super().update(inner_aval, memory_space, kind)
    return AbstractRefUnion(ref.inner_aval, self.refs, self.memory_space)

  @functools.cached_property
  def layout(self) -> tcgen05.TMEMLayout:
    if self.memory_space != TMEM:
      raise ValueError("layout attribute is only defined for TMEM refs")
    return tcgen05.tmem_default_layout(packing=1)

  @functools.cached_property
  def collective(self) -> bool:
    if self.memory_space != TMEM:
      raise ValueError("collective attribute is only defined for TMEM refs")
    ref_leaves = jax.tree.leaves(self.refs)
    first_ref = ref_leaves[0]
    assert all(ref.collective == first_ref.collective for ref in ref_leaves)
    return first_ref.collective

  def __eq__(self, other):
    return (
        type(self) is type(other)
        and self.inner_aval == other.inner_aval
        and self.memory_space == other.memory_space
        and self.refs == other.refs
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check `ref.memory_space == TMEM` before accessing `.layout` (or use `getattr(ref, 'memory_space', None)` dispatch)
2. Branch on memory_space and only query layout for TMEM refs
3. If you expected TMEM, verify the ref was actually allocated in TMEM (e.g. via `pl.tmem_ref` / tcgen05 APIs) rather than SMEM

Example fix

# before
layout = ref_aval.layout  # ValueError for SMEM unions

# after
if ref_aval.memory_space == mosaic_gpu_core.TMEM:
    layout = ref_aval.layout
else:
    layout = None  # SMEM refs have no layout

if ref_aval.memory_space == mosaic_gpu_core.TMEM:
    layout = ref_aval.layout
else:
    layout = None

def is_tmem_ref_aval(a) -> bool:
    return getattr(a, 'memory_space', None) == mosaic_gpu_core.TMEM

null

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• collective attribute is only defined for TMEM refs
• Some aliased TMEM references are collective and some are not
• All aliased Refs must have the same memory space (SMEM or TM

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