jax-ml/jax · error · ValueError

Ref unions can't be assigned to.

Error message

Ref unions can't be assigned to.

What it means

A Mosaic GPU ref union (`AbstractRefUnion`) is a read-only container of several SMEM/TMEM refs grouped for layout purposes. JAX's abstract-ref protocol asks for `_setitem` when the tracer is assigned to; the union deliberately raises ValueError because assignment to the union as a whole is meaningless.

Source

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

  def __init__(
      self,
      aval,
      refs: Sequence[_GPUMemoryRefTree],
      memory_space,
  ):
    self.refs = refs
    super().__init__(aval, memory_space=memory_space)

  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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Iterate the union (`for r in ref_union: ...`) and assign to each member ref individually
2. Restructure the kernel to keep member refs separately and only use the union where a grouped layout is required
3. If a JAX transform internally triggers setitem, avoid passing the union through that transform (e.g. don't vmap over it)

Example fix

# before
ref_union[0] = value  # ValueError: Ref unions can't be assigned to.
# after
for r in ref_union:
    r[...] = value

def is_ref_union(x) -> bool:
    from jax._src.pallas.mosaic_gpu.core import AbstractRefUnion
    return isinstance(getattr(x, 'aval', x), AbstractRefUnion)

try:
    ref_union[i] = v
except ValueError:
    for r in ref_union:
        r[...] = v

When it happens

Trigger: Attempting `ref_union[idx] = value` or any indexed assignment on a tracer whose aval is AbstractRefUnion (e.g. inside a Pallas kernel body or under vmap/jaxpr tracing that performs setitem on the union).

Common situations: Writing to scratch memory via the union object instead of iterating its member refs; autodiff/vmap machinery or custom rules trying to update the union in place; treating the union like a normal BlockRef.

Related errors

• Only SMEM and TMEM refs are supported.
• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• Cannot specify both out_shape and out_type. Use out_type.
• Cannot specify both scratch_shapes and scratch_types. Use sc

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