jax-ml/jax · error · ValueError

Expected metadata dtype to be uint2, got: {meta.dtype}

Error message

Expected metadata dtype to be uint2, got: {meta.dtype}

What it means

format_tcgen05_sparse_metadata requires the sparse metadata array to have dtype uint2 (the tcgen05 sparse MMA metadata format). Any other dtype is rejected before the reshaping/tiling logic.

Source

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

  return decorator


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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast before calling: meta = meta.astype(jax.numpy.uint2)
2. Generate metadata directly as uint2 (e.g. via arange % 4 cast to uint2 in tests)

Example fix

// before
meta = meta.astype(jnp.uint8)
formatted = format_tcgen05_sparse_metadata(meta, jnp.float8_e4m3fn)

// after
meta = meta.astype(jnp.uint2)
formatted = format_tcgen05_sparse_metadata(meta, jnp.float8_e4m3fn)

import jax.numpy as jnp
if meta.dtype != jnp.uint2:
    meta = meta.astype(jnp.uint2)

def metadata_dtype_ok(meta) -> bool:
    return str(meta.dtype) == 'uint2'

try:
    return format_tcgen05_sparse_metadata(meta, dt)
except ValueError:
    return format_tcgen05_sparse_metadata(meta.astype(jnp.uint2), dt)

When it happens

Trigger: Passing metadata as uint8/uint16/int32 to format_tcgen05_sparse_metadata(meta, operand_dtype) — e.g. metadata loaded from a file or produced by np.packbits without casting.

Common situations: Preparing sparse metadata outside JAX (numpy defaults to uint8/uint64) and forgetting the cast; porting from a different sparse format where metadata was wider.

Related errors

• Sparse metadata format not implemented for {operand_dtype=}
• Expected metadata to be 3-dimensional (M, K // 4, 2), but it
• Expected the trailing dimension of the metadata to be 2, got
• Unsupported dtype: {ref.dtype}
• dims and idxs must have the same length

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