jax-ml/jax · error · ValueError

Expected the trailing dimension of the metadata to be 2, got

Error message

Expected the trailing dimension of the metadata to be 2, got: {meta.shape[-1]}

What it means

Within the required 3D (M, K // 4, 2) metadata layout, the trailing dimension must be exactly 2. A different last-dim size means the metadata was packed incorrectly for the tcgen05 sparse format.

Source

Thrown at jax/_src/pallas/mosaic_gpu/helpers.py:209

  for the documentation of the required layouts. The array can be copied into
  SMEM, from where ``plgpu.async_copy_sparse_metadata_to_tmem`` can be used to
  copy it over to TMEM. The formatting of the array depends on the data type of
  the operands to the sparse MMA operation.

  Args:
    meta: Metadata of shape (M, K // 4, 2).
    dtype: Data type of MMA operands.
  """
  if meta.dtype != dtypes.uint2:
    raise ValueError(f"Expected metadata dtype to be uint2, got: {meta.dtype}")
  if meta.ndim != 3:
    raise ValueError(
        "Expected metadata to be 3-dimensional (M, K // 4, 2), but it is"
        f" {meta.ndim}D"
    )
  m, k, _2 = meta.shape
  if _2 != 2:
    raise ValueError(
        "Expected the trailing dimension of the metadata to be 2, got:"
        f" {meta.shape[-1]}"
    )
  k *= 2
  bitsize = dtypes.itemsize_bits(operand_dtype)
  if bitsize == 8:
    meta_tiled = meta.reshape(m // 128, 128, k // 64, 64).transpose(0, 2, 1, 3)
  elif bitsize == 16:
    meta_tiled = meta.reshape(m // 128, 8, 2, 8, k // 64, 4, 2, 8).transpose(0, 4, 1, 6, 3, 5, 2, 7)
  else:
    raise NotImplementedError(
        f"Sparse metadata format not implemented for {operand_dtype=}"
    )
  return meta_tiled.reshape(m // 128, k // 64, 128, 64)


def find_swizzle(minor_dim_bits: int, what: str = ""):
  """Returns the largest swizzle that can be applied to a memory region.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Fix the packing so the final axis holds the 2 metadata entries: meta.reshape(m, k // 4, 2)
2. Double-check against the docstring: shape (M, K // 4, 2), dtype uint2

Example fix

// before
meta = meta.reshape(m, k // 2, 1)

// after
meta = meta.reshape(m, k // 4, 2)

assert meta.shape[-1] == 2, f'trailing dim must be 2, got {meta.shape[-1]}'

def trailing_dim_is_2(meta) -> bool:
    return meta.ndim >= 1 and meta.shape[-1] == 2

try:
    return format_tcgen05_sparse_metadata(meta, dt)
except ValueError:
    meta = meta.reshape(*meta.shape[:-1], -1)
    meta = meta.reshape(*meta.shape[:-1], meta.shape[-1] // 2, 2)
    return format_tcgen05_sparse_metadata(meta, dt)

When it happens

Trigger: meta.shape[-1] != 2, e.g. reshaping to (M, K//4, 4) or (M, K//2, 1) before calling format_tcgen05_sparse_metadata.

Common situations: Mistaking the layout as (M, K//4, 1) pairs or interleaving metadata differently when preprocessing weights for sparse MMA.

Related errors

• Expected metadata to be 3-dimensional (M, K // 4, 2), but it
• Expected metadata dtype to be uint2, got: {meta.dtype}
• Sparse metadata format not implemented for {operand_dtype=}
• Stores to TMEM are asynchronous operations and cannot be per
• MMA rhs swizzle must match lhs swizzle. {lhs_swizzle=} {rhs_

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