{"record":{"id":"32f0da841bf70011","repo":"jax-ml/jax","slug":"ref-unions-can-t-be-assigned-to","errorCode":null,"errorMessage":"Ref unions can't be assigned to.","messagePattern":"Ref unions can't be assigned to\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":619,"sourceCode":"\n  def __init__(\n      self,\n      aval,\n      refs: Sequence[_GPUMemoryRefTree],\n      memory_space,\n  ):\n    self.refs = refs\n    super().__init__(aval, memory_space=memory_space)\n\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)","sourceCodeStart":601,"sourceCodeEnd":637,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L601-L637","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Iterate the union (`for r in ref_union: ...`) and assign to each member ref individually","Restructure the kernel to keep member refs separately and only use the union where a grouped layout is required","If a JAX transform internally triggers setitem, avoid passing the union through that transform (e.g. don't vmap over it)"],"exampleFix":"# before\nref_union[0] = value  # ValueError: Ref unions can't be assigned to.\n# after\nfor r in ref_union:\n    r[...] = value","handlingStrategy":"fallback","validationCode":null,"typeGuard":"def is_ref_union(x) -> bool:\n    from jax._src.pallas.mosaic_gpu.core import AbstractRefUnion\n    return isinstance(getattr(x, 'aval', x), AbstractRefUnion)","tryCatchPattern":"try:\n    ref_union[i] = v\nexcept ValueError:\n    for r in ref_union:\n        r[...] = v","preventionTips":["Always write through member refs, never the union","Keep member refs accessible alongside any union you build"],"tags":["jax","pallas","mosaic-gpu","ref-union","immutable","read-only"],"backgroundTag":"readonly-container-assignment","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}