{"record":{"id":"4c248a06e138ecf3","repo":"jax-ml/jax","slug":"packed-cannot-be-specified-if-layout-is-specified","errorCode":null,"errorMessage":"packed cannot be specified if layout is specified.","messagePattern":"packed cannot be specified if layout is specified\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":201,"sourceCode":"      if collective is None:\n        collective = False\n      if len(shape) > 2:\n        transforms = (CollapseLeadingBatchDimensionsTransform(),)\n      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, ...]):","sourceCodeStart":183,"sourceCodeEnd":219,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L183-L219","documentation":"In JAX's Mosaic GPU (Pallas) API, a TMEM (tensor memory) ref can be allocated with either an explicit `layout` or with `packed`/`collective` hints from which the layout is inferred — not both. When you call a TMEM memory space's ref constructor with `layout` set and also pass `packed`, this ValueError fires because packedness would conflict with (or be redundant to) the explicitly given layout.","triggerScenarios":"Calling `tmem_allocator.get_buf(...)` or `MemorySpace.__call__`/`like` on a TMEM space with both `layout=...` (e.g. `TMEMLayout(...)` or mgpu layout) and `packed=True/False` specified. Non-None `packed` while `layout is not None` hits the raise.","commonSituations":"Upgrading Pallas kernels that previously passed `packed=True` for TMEM and then adding an explicit `layout` argument (newer JAX added layout support); copy-pasting a TMEM allocation that combined options from two examples.","solutions":["Remove the `packed` argument and let the explicit `layout` fully define the TMEM layout","Or drop `layout` and keep `packed`/`collective` so the layout is inferred via infer_tmem_layout","Check the JAX version changelog — `layout` on TMEM refs is a newer API; on old versions only packed/collective exist"],"exampleFix":"# before\nbuf = smem_tmem.get_buf(shape=(128, 128), dtype=jnp.float32, packed=True, layout=TMEMLayout(...))\n# after\nbuf = smem_tmem.get_buf(shape=(128, 128), dtype=jnp.float32, layout=TMEMLayout(...))","handlingStrategy":"validation","validationCode":"def alloc_tmem(space, *, layout=None, packed=None, collective=None, **kw):\n    if layout is not None and packed is not None:\n        raise ValueError('pass either layout or packed, not both')\n    return space.get_buf(layout=layout, packed=packed, collective=collective, **kw)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Centralize TMEM allocation in one helper that rejects conflicting kwargs","Keep packed/collective only on allocations without an explicit layout"],"tags":["jax","pallas","mosaic-gpu","tmem","layout","api-misuse"],"backgroundTag":"mutually-exclusive-arguments","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}