jax-ml/jax · error · ValueError

A sparse metadata shape mismatch: expected {(m, expected_met

Error message

A sparse metadata shape mismatch: expected {(m, expected_meta_k)}, got {a_sparse_metadata.shape}

What it means

In sparse MMA, the A sparse metadata must have shape (m, k // sparse_group_elems * 2), where sparse_group_elems is 8 for 4-bit and 4 for other types (each sparse group stores 2 of the entries). This error fires when the metadata shape deviates.

Source

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

        )
    else:
      raise AssertionError("Should not happen")
    if b_scale.shape[0] % 128 or b_scale.shape[0] < n * num_cta:
      raise ValueError(
          f"B scale shape[0] must be a multiple of 128 and >= N={n * num_cta},"
          f" got {b_scale.shape[0]}"
      )
    if b_scale.shape[1] != k_scales:
      raise ValueError(
          f"B scale shape mismatch: expected ({b_scale.shape[0]}, {k_scales}),"
          f" got {b_scale.shape}"
      )
  if is_sparse:
    sparse_group_elems = 8 if utils.bitwidth(a_element_type) == 4 else 4
    # Each sparse group has 2 entries.
    expected_meta_k = k // sparse_group_elems * 2
    if a_sparse_metadata.shape != (m, expected_meta_k):
      raise ValueError(
          f"A sparse metadata shape mismatch: expected {(m, expected_meta_k)},"
          f" got {a_sparse_metadata.shape}"
      )
    if a_sparse_metadata.dtype != ir.IntegerType.get_signless(2):
      raise ValueError(
          "A sparse metadata dtype mismatch: expected i2, got"
          f" {a_sparse_metadata.dtype}"
      )

  # Step 3. Compute the operand descriptors.
  if not isinstance(a, TMEMRef):
    # Both dense and sparse matmul consume A with a K bytewidth of 32, only
    # the group size is halved when it's sparse.
    (
        (a_desc_base, a_k_instr_strides),
        (a_m_group_stride, a_k_group_stride),
        a_fastest,
    ) = mma_utils.create_descriptor(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Compute metadata as shape (m, k // (8 if bitwidth==4 else 4) * 2)
2. Ensure the metadata M matches the MMA m dimension exactly

Example fix

# before
meta = build_metadata((m, k))
# after
sparse_group_elems = 8 if bitwidth(a_dtype) == 4 else 4
meta = build_metadata((m, k // sparse_group_elems * 2))

sge = 8 if utils.bitwidth(a_element_type) == 4 else 4
assert a_sparse_metadata.shape == (m, k // sge * 2)

def valid_sparse_meta(shape, m, k, dtype) -> bool:
    sge = 8 if utils.bitwidth(dtype) == 4 else 4
    return tuple(shape) == (m, k // sge * 2)

When it happens

Trigger: Passing a_sparse_metadata with the wrong M or K-derived dimension, e.g. using k instead of k//4*2 for an 8-bit type, or M not matching the a tile.

Common situations: Metadata generated with the wrong E2EM (2:4 sparsity) group size after changing operand dtype between f8 and f4; M-dim mismatch when a is a TMEMRef slice.

Related errors

• A sparse metadata dtype mismatch: expected i2, got {a_sparse
• A scale shape mismatch: expected ({TMEM_ROWS}, {k_scales}),
• B scale shape[0] must be a multiple of 128 and >= N={n * num
• B scale shape mismatch: expected ({b_scale.shape[0]}, {k_sca
• B tiling too small. Increase swizzle or transpose the input.

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