{"record":{"id":"166e4bcf455e6acb","repo":"jax-ml/jax","slug":"packed-collective-and-layout-arguments-are-only-s","errorCode":null,"errorMessage":"packed, collective and layout arguments are only supported for TMEM.","messagePattern":"packed, collective and layout arguments are only supported for TMEM\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":205,"sourceCode":"      if layout is None:\n        if packed is None:\n          if dtypes.itemsize_bits(dtype) != 32:\n            raise ValueError(\n                \"dtypes narrower than 32-bit require either the packed argument\"\n                \" or an explicit TMEM layout\"\n            )\n          packed = False\n        # Ignore batch dimensions for layout inference.\n        mgpu_layout = infer_tmem_layout(\n            shape[-2:], dtype, packed=packed, collective=collective\n        )\n      else:\n        if packed is not None:\n          raise ValueError(\"packed cannot be specified if layout is specified.\")\n        mgpu_layout = layout.to_mgpu()\n    else:\n      if packed is not None or collective is not None or layout is not None:\n        raise ValueError(\"packed, collective and layout arguments are only supported for TMEM.\")\n      mgpu_layout = None\n    return GPUMemoryRef(jax_core.ShapedArray(shape, dtype), memory_space=self,\n                        transforms=transforms, layout=mgpu_layout,\n                        collective=collective)\n\n  def like(self, shape_dtype_like):\n    return self(shape_dtype_like.shape, shape_dtype_like.dtype)\n\n\nclass SemaphoreType(enum.Enum):\n  REGULAR = \"regular\"\n  BARRIER = \"barrier\"\n\n  def __call__(self, shape: tuple[int, ...]):\n    dtype: Any\n    if self == SemaphoreType.BARRIER:\n      dtype = pallas_core.BarrierSemaphore()\n    else:","sourceCodeStart":187,"sourceCodeEnd":223,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L187-L223","documentation":"The `packed`, `collective`, and `layout` keyword arguments of the Pallas Mosaic GPU memory-space ref constructor only make sense for TMEM (tensor core memory). Passing any of them when allocating a ref in another memory space (SMEM, gmem) is rejected with this ValueError.","triggerScenarios":"Allocating an SMEM or GMEM ref, e.g. `smem.get_buf(..., packed=True)` or `grid_mem.get_ref(..., collective=...)`, where the memory space is not TMEM, so any of the three kwargs being non-None triggers the raise.","commonSituations":"Reusing a TMEM allocation call for an SMEM scratch buffer during kernel refactoring; writing hardware-agnostic Pallas code and forwarding the same kwargs to all memory spaces; version changes where `collective` was introduced and older tutorial code spreads it everywhere.","solutions":["Remove `packed`/`collective`/`layout` from the non-TMEM ref allocation","Move those kwargs so they only apply to the TMEM allocator call","If you intended TMEM semantics, make sure you are calling the TMEM memory space (e.g. the tensor-memory allocator), not SMEM/GMEM"],"exampleFix":"# before\nsmem_buf = smem.get_buf(shape=(64, 64), dtype=jnp.float32, packed=True)\n# after\nsmem_buf = smem.get_buf(shape=(64, 64), dtype=jnp.float32)","handlingStrategy":"validation","validationCode":"def get_buf(space, **kw):\n    if not getattr(space, 'is_tmem', lambda: False)():\n        kw.pop('packed', None); kw.pop('collective', None); kw.pop('layout', None)\n    return space(**kw)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Only forward packed/collective/layout to TMEM allocators","Assert memory space type before passing TMEM-only kwargs"],"tags":["jax","pallas","mosaic-gpu","smem","gmem","tmem","api-misuse"],"backgroundTag":"invalid-argument-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}