jax-ml/jax · error · ValueError

packed, collective and layout arguments are only supported f

Error message

packed, collective and layout arguments are only supported for TMEM.

What it means

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.

Source

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

      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, ...]):
    dtype: Any
    if self == SemaphoreType.BARRIER:
      dtype = pallas_core.BarrierSemaphore()
    else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove `packed`/`collective`/`layout` from the non-TMEM ref allocation
2. Move those kwargs so they only apply to the TMEM allocator call
3. If you intended TMEM semantics, make sure you are calling the TMEM memory space (e.g. the tensor-memory allocator), not SMEM/GMEM

Example fix

# before
smem_buf = smem.get_buf(shape=(64, 64), dtype=jnp.float32, packed=True)
# after
smem_buf = smem.get_buf(shape=(64, 64), dtype=jnp.float32)

def get_buf(space, **kw):
    if not getattr(space, 'is_tmem', lambda: False)():
        kw.pop('packed', None); kw.pop('collective', None); kw.pop('layout', None)
    return space(**kw)

When it happens

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

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

Related errors

• packed cannot be specified if layout is specified.
• All aliased Refs must have the same memory space (SMEM or TM
• 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/166e4bcf455e6acb. Report an issue: GitHub.