{"record":{"id":"70a49296e69993ec","repo":"jax-ml/jax","slug":"only-smem-and-tmem-refs-are-supported","errorCode":null,"errorMessage":"Only SMEM and TMEM refs are supported.","messagePattern":"Only SMEM and TMEM refs are supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":595,"sourceCode":"      def unflatten(ref):\n        nonlocal col_offset\n        col_offset = align_to(col_offset, TMEM_COL_ALIGNMENT)\n        if not isinstance(ref, pallas_core.TransformedRef):\n          ref = pallas_core.TransformedRef(ref, transforms=())\n        ncols = ref.layout.cols_in_shape(ref.shape,\n                                         dtypes.itemsize_bits(ref.dtype))\n        transform = ExtractAliasedRef.from_transformed_ref(\n            ref, col_offset, group_idx, layout=ref.layout)\n        result = pallas_core.TransformedRef(\n            ref_union, transforms=(transform, *ref.transforms)\n        )\n        col_offset += ncols\n        return result\n      flat_refs.append(jax.tree.map(unflatten, ref_group))\n      union_cols = max(union_cols, col_offset)\n    assert union_cols == ref_union.shape[1], (union_cols, ref_union.shape[1])\n  else:\n    raise NotImplementedError(\"Only SMEM and TMEM refs are supported.\")\n  return tuple(flat_refs)\n\n\nclass AbstractRefUnion(state.AbstractRef):\n  refs: Sequence[_GPUMemoryRefTree]\n\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","sourceCodeStart":577,"sourceCodeEnd":613,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L577-L613","documentation":"`flatten_ref_union` only knows how to lay out refs living in shared memory (SMEM) or tensor memory (TMEM). Calling it on refs whose memory space is anything else — typically global memory (GMEM) refs or host/abstract refs — raises NotImplementedError.","triggerScenarios":"Passing a grid/gmem BlockRef (an input or output of a Pallas kernel) into `flatten_ref_union`, or any ref tree containing a ref whose memory_space is not SMEM or TMEM.","commonSituations":"Accidentally feeding kernel I/O refs (which live in GMEM) into union-building helpers meant for scratch buffers; generic tree-mapping code that hits both scratch and I/O refs; assuming the helper is memory-space agnostic.","solutions":["Filter the ref tree so only SMEM/TMEM scratch refs reach flatten_ref_union","Keep GMEM input/output refs out of the union; process them separately in the kernel","Check ref.memory_space before grouping if refs come from mixed sources"],"exampleFix":"# before\nflat = flatten_ref_union((x_ref, smem_scratch))  # x_ref is a GMEM input\n# after\nflat = flatten_ref_union((smem_scratch,))  # only SMEM/TMEM refs","handlingStrategy":"type-guard","validationCode":"def only_onchip(refs):\n    return [r for r in refs if type(r.memory_space).__name__ in ('SMEMAddressSpace','TMEMAddressSpace') or str(r.memory_space) in ('SMEM','TMEM')]","typeGuard":"def is_onchip_ref(ref) -> bool:\n    ms = getattr(ref, 'memory_space', None)\n    return ms is not None and str(getattr(ms, 'memory_space', ms)) in ('SMEM', 'TMEM')","tryCatchPattern":null,"preventionTips":["Filter kernel I/O (GMEM) refs out before calling flatten_ref_union","Keep scratch-ref grouping code separate from input/output handling"],"tags":["jax","pallas","mosaic-gpu","ref-union","memory-space","not-implemented"],"backgroundTag":"unsupported-memory-space","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}