jax-ml/jax · error · ValueError

Packing must be a power of 2, got: {packing}

Error message

Packing must be a power of 2, got: {packing}

What it means

tmem_default_layout builds the standard 128-row TMEM layout with a packing (vector width along columns) that must be a power of 2 because it maps directly onto vectorized TMEM loads/stores. packing.bit_count() != 1 catches zero, negatives, and non-powers like 3 or 6.

Source

Thrown at jax/experimental/mosaic/gpu/tcgen05.py:1087

    raise ValueError(f"Minor dimension of shape must be divisible by packing, got: {shape}")
  if shape[0] == TMEM_ROWS:
    return tmem_default_layout(packing)
  elif shape[0] == TMEM_ROWS // 2:
    if collective:
      return tmem_m64_collective_layout(shape[1], packing)
    else:
      return tmem_half_lane_layout(shape[1], packing)
  else:
    raise ValueError(
        f"Unsupported shape: {shape}. TMEM references must have either"
        f" {TMEM_ROWS} or {TMEM_ROWS // 2} rows, but got {shape[0]}."
    )


def tmem_default_layout(packing: int = 1) -> TMEMLayout:
  """A TMEM layout used for 1CTA MMA with M=128 and 2CTA MMA with M=256."""
  if packing.bit_count() != 1:
    raise ValueError(f"Packing must be a power of 2, got: {packing}")
  return TMEMLayout(
      fa.Tiling(((TMEM_ROWS, packing), (fa.WARP_SIZE, packing))),
      warp_dims=(-4,),
      lane_dims=(-2,),
      vector_dim=-1,
  )


def tmem_half_lane_layout(columns, packing: int = 1) -> TMEMLayout:
  """A TMEM layout used for 1CTA MMA with M=64."""
  if packing > (columns // 2) or packing.bit_count() != 1:
    raise ValueError(f"Packing must be <= 8 and a power of 2, got: {packing}")
  if columns % 16:
    raise ValueError(f"Columns must be a multiple of 16, got: {columns}")
  return TMEMLayout(
      fa.Tiling((
          (TMEM_ROWS // 2, columns),
          (fa.WARP_SIZE // 2, columns // 2),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a power of 2: 1, 2, 4, or 8
2. If packing comes from bitwidth math, clamp/round down to the nearest power of two
3. For packing=0, default the parameter instead of passing it explicitly

Example fix

# before
layout = tcgen05.tmem_default_layout(packing=3)
# after
layout = tcgen05.tmem_default_layout(packing=2)

assert isinstance(packing, int) and packing > 0 and packing.bit_count() == 1, packing

def is_pow2(n: int) -> bool:
    return isinstance(n, int) and n > 0 and n.bit_count() == 1

When it happens

Trigger: Calling tmem_default_layout(3), tmem_default_layout(0), or passing a computed packing such as bitwidth-derived values that aren't powers of two; also reached indirectly via layout-related helpers (is_valid_tmem_transfer, pprint_layout, async store constraint system).

Common situations: Deriving packing from element bitwidth with arithmetic that yields non-power-of-2 (e.g. 24/8=3 is fine but 48/32 combinations can yield 6); copying example code with packing hardcoded incorrectly.

Understand the failure class

Background: "Must be a positive integer", "Invalid value", "Unsupported": the invalid-argument-value error family, when a library rejects the value you pass — this error's family across 35 libraries.

Related errors

• Minor dimension of shape must be divisible by packing, got:
• Columns must be a multiple of 16, got: {columns}
• tmem_addr_ref must be a memref or a pointer, got: {tmem_addr
• Cannot slice TMEM with multiple tiles along rows.
• TMEM layout {self.layout} is not supported

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