jax-ml/jax · error · ValueError

Expected metadata to be 3-dimensional (M, K // 4, 2), but it

Error message

Expected metadata to be 3-dimensional (M, K // 4, 2), but it is {meta.ndim}D

What it means

format_tcgen05_sparse_metadata expects metadata shaped (M, K // 4, 2) — exactly 3D. Higher- or lower-rank arrays are rejected because the subsequent reshape/transpose logic assumes this exact structure.

Source

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

def format_tcgen05_sparse_metadata(meta, operand_dtype):
  """Formats the sparse metadata for tcgen05.mma into the expected format.

  See
  https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-sparse-matrices-sparsity-selector-kind-f16-m128-256
  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=}"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Reshape to (M, K // 4, 2) before the call, e.g. meta.reshape(m, k // 4, 2) after ensuring dtype uint2
Remove accidental batch axes (squeeze) introduced by vmap or stacking

Example fix

// before
meta = packed.reshape(m, k // 2)  # 2D, wrong

// after
meta = packed.reshape(m, k // 4, 2).astype(jnp.uint2)

Defensive patterns

Strategy: validation

Validate before calling

assert meta.ndim == 3, f'expected (M, K//4, 2), got {meta.shape}'

Type guard

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

Try / catch

try:
    return format_tcgen05_sparse_metadata(meta, dt)
except ValueError:
    meta = meta.reshape(meta.shape[0], -1 // 2 if meta.ndim == 2 else meta.shape[1], 2)
    return format_tcgen05_sparse_metadata(meta.astype(jnp.uint2), dt)

Prevention

Construct metadata as (M, K//4, 2) from the start
squeeze batch axes introduced by vmap before formatting

When it happens

Trigger: Passing a 2D packed metadata array or a 4D batched array to format_tcgen05_sparse_metadata.

Common situations: Pre-packed/flattened metadata from a data pipeline; vmap/batching accidentally adding a leading dimension; misconverting the (M, K//4, 2) layout from documentation.

Related errors

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