jax-ml/jax · error · ValueError

collective attribute is only defined for TMEM refs

Error message

collective attribute is only defined for TMEM refs

What it means

AbstractRefUnion.collective in Mosaic GPU core returns whether the aliased TMEM references use the tcgen05 collective (warpgroup-wide) access mode. It is only defined for TMEM unions; accessing it on SMEM refs raises ValueError. The value is taken from the first leaf ref, asserting all leaves agree.

Source

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

  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
    )

  def __hash__(self):
    # `flatten_ref_union(self)` creates `TransformedRef`s that refer to `self`,
    # so we extract the transforms from the `TransformedRef`s in order to avoid
    # infinite recursion.
    all_transforms = tuple(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Guard with `ref.memory_space == TMEM` before reading `.collective`
2. When creating TMEM refs, specify collective explicitly (tcgen05 allocation API) so the attribute is well-defined
3. Audit code paths that assume collective semantics and split TMEM vs SMEM handling

Example fix

# before
is_collective = ref_aval.collective

# after
is_collective = (
    ref_aval.collective
    if ref_aval.memory_space == mosaic_gpu_core.TMEM
    else False
)

is_collective = (
    ref_aval.collective
    if ref_aval.memory_space == mosaic_gpu_core.TMEM
    else False
)

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

null

When it happens

Trigger: Accessing `.collective` on an AbstractRefUnion with memory_space != TMEM, e.g. when interpreting or lowering kernels that mix SMEM and TMEM refs (tests like testArgAllReduce / testCommAssocCollective exercise the TMEM path).

Common situations: Generic introspection over ref avals during tracing/interpretation; refactors that pass SMEM refs into code paths written for TMEM collective ops; JAX version upgrades where the collective attribute was added to the API surface.

Related errors

• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• layout 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/8c4fa3749030da19. Report an issue: GitHub.