jax-ml/jax · error · ValueError

packed cannot be specified if layout is specified.

Error message

packed cannot be specified if layout is specified.

What it means

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.

Source

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

      if collective is None:
        collective = False
      if len(shape) > 2:
        transforms = (CollapseLeadingBatchDimensionsTransform(),)
      if layout is None:
        if packed is None:
          if dtypes.itemsize_bits(dtype) != 32:
            raise ValueError(
                "dtypes narrower than 32-bit require either the packed argument"
                " or an explicit TMEM layout"
            )
          packed = False
        # Ignore batch dimensions for layout inference.
        mgpu_layout = infer_tmem_layout(
            shape[-2:], dtype, packed=packed, collective=collective
        )
      else:
        if packed is not None:
          raise ValueError("packed cannot be specified if layout is specified.")
        mgpu_layout = layout.to_mgpu()
    else:
      if packed is not None or collective is not None or layout is not None:
        raise ValueError("packed, collective and layout arguments are only supported for TMEM.")
      mgpu_layout = None
    return GPUMemoryRef(jax_core.ShapedArray(shape, dtype), memory_space=self,
                        transforms=transforms, layout=mgpu_layout,
                        collective=collective)

  def like(self, shape_dtype_like):
    return self(shape_dtype_like.shape, shape_dtype_like.dtype)


class SemaphoreType(enum.Enum):
  REGULAR = "regular"
  BARRIER = "barrier"

  def __call__(self, shape: tuple[int, ...]):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the `packed` argument and let the explicit `layout` fully define the TMEM layout
2. Or drop `layout` and keep `packed`/`collective` so the layout is inferred via infer_tmem_layout
3. Check the JAX version changelog — `layout` on TMEM refs is a newer API; on old versions only packed/collective exist

Example fix

# before
buf = smem_tmem.get_buf(shape=(128, 128), dtype=jnp.float32, packed=True, layout=TMEMLayout(...))
# after
buf = smem_tmem.get_buf(shape=(128, 128), dtype=jnp.float32, layout=TMEMLayout(...))

def alloc_tmem(space, *, layout=None, packed=None, collective=None, **kw):
    if layout is not None and packed is not None:
        raise ValueError('pass either layout or packed, not both')
    return space.get_buf(layout=layout, packed=packed, collective=collective, **kw)

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• packed, collective and layout arguments are only supported f
• Can't instantiate {self} with arguments.
• layout attribute is only defined for TMEM refs
• collective attribute is only defined for TMEM refs
• Some aliased TMEM references are collective and some are not

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