jax-ml/jax · error · ValueError

Minor dimension of shape must be divisible by packing, got:

Error message

Minor dimension of shape must be divisible by packing, got: {shape}

What it means

Tensor Memory (TMEM) on Blackwell GPUs is addressed in packed columns; the inferred TMEM layout requires the minor (column) dimension of the 2D shape to be divisible by the packing factor. _infer_tmem_layout validates this before constructing a TMEMLayout. If shape[1] % packing != 0, no valid layout exists, so the error is raised.

Source

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

        layout.warp_dims,
        layout.lane_dims,
        layout.vector_dim,
        _check_canonical=False,
    )

  def as_tiled_layout(self) -> fa.TiledLayout:
    return fa.TiledLayout(
        self.tiling, self.warp_dims, self.lane_dims, self.vector_dim
    )


def _infer_tmem_layout(shape: tuple[int, ...], collective: bool, packing: int) -> TMEMLayout:
  if len(shape) != 2:
    raise ValueError(f"TMEM can only represent 2D shapes, got {shape}")
  if packing > 8 or packing.bit_count() != 1:
    raise ValueError(f"Packing must be <= 8 and a power of 2, got: {packing}")
  if shape[1] % packing:
    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}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Pad the minor dimension of the shape up to the next multiple of packing (e.g. N=10 with packing=4 -> N=12)
  2. Reduce the packing factor so it divides the column count (packing=1 always works)
  3. Pass an explicit layout= to from_alloc instead of relying on inference

Example fix

# before
ref = tcgen05.TMEMRef.from_alloc(alloc, (128, 10), collective=True, packing=4)
# after
ref = tcgen05.TMEMRef.from_alloc(alloc, (128, 12), collective=True, packing=4)  # pad N to multiple of 4
Defensive patterns

Strategy: validation

Validate before calling

def check_tmem_shape(shape, packing):
    assert len(shape) == 2 and shape[1] % packing == 0, f'{shape} not divisible by packing={packing}'

Type guard

def has_valid_packing(shape: tuple[int, ...], packing: int) -> bool:
    return len(shape) == 2 and shape[1] % packing == 0

Try / catch

try:
    ref = tcgen05.TMEMRef.from_alloc(alloc, shape, collective=c)
except ValueError as e:
    if 'divisible by packing' in str(e):
        shape = (shape[0], (shape[1] + packing - 1) // packing * packing)
    else:
        raise

Prevention

When it happens

Trigger: Calling infer_tmem_layout / from_alloc (with layout=None) or mma with a shape whose second dimension is not a multiple of the packing (e.g. packing=4 with shape=(128, 10)). Also triggered from _construct_smem_reftree or _default_tmem_layout_for_variable when a variable's column count is odd relative to packing.

Common situations: Mosaic GPU kernels using tcgen05 MMA where the N dimension of the accumulator is not padded to a power-of-2 multiple; using packing inferred from element bitwidth (e.g. packing=8 for i4/f8) with an N that isn't divisible by 8.

Related errors

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