{"record":{"id":"2e3e5acb0313dc4e","repo":"jax-ml/jax","slug":"a-sparse-metadata-shape-mismatch-expected-m-ex","errorCode":null,"errorMessage":"A sparse metadata shape mismatch: expected {(m, expected_meta_k)}, got {a_sparse_metadata.shape}","messagePattern":"A sparse metadata shape mismatch: expected (.+?), got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":501,"sourceCode":"        )\n    else:\n      raise AssertionError(\"Should not happen\")\n    if b_scale.shape[0] % 128 or b_scale.shape[0] < n * num_cta:\n      raise ValueError(\n          f\"B scale shape[0] must be a multiple of 128 and >= N={n * num_cta},\"\n          f\" got {b_scale.shape[0]}\"\n      )\n    if b_scale.shape[1] != k_scales:\n      raise ValueError(\n          f\"B scale shape mismatch: expected ({b_scale.shape[0]}, {k_scales}),\"\n          f\" got {b_scale.shape}\"\n      )\n  if is_sparse:\n    sparse_group_elems = 8 if utils.bitwidth(a_element_type) == 4 else 4\n    # Each sparse group has 2 entries.\n    expected_meta_k = k // sparse_group_elems * 2\n    if a_sparse_metadata.shape != (m, expected_meta_k):\n      raise ValueError(\n          f\"A sparse metadata shape mismatch: expected {(m, expected_meta_k)},\"\n          f\" got {a_sparse_metadata.shape}\"\n      )\n    if a_sparse_metadata.dtype != ir.IntegerType.get_signless(2):\n      raise ValueError(\n          \"A sparse metadata dtype mismatch: expected i2, got\"\n          f\" {a_sparse_metadata.dtype}\"\n      )\n\n  # Step 3. Compute the operand descriptors.\n  if not isinstance(a, TMEMRef):\n    # Both dense and sparse matmul consume A with a K bytewidth of 32, only\n    # the group size is halved when it's sparse.\n    (\n        (a_desc_base, a_k_instr_strides),\n        (a_m_group_stride, a_k_group_stride),\n        a_fastest,\n    ) = mma_utils.create_descriptor(","sourceCodeStart":483,"sourceCodeEnd":519,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L483-L519","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Compute metadata as shape (m, k // (8 if bitwidth==4 else 4) * 2)","Ensure the metadata M matches the MMA m dimension exactly"],"exampleFix":"# before\nmeta = build_metadata((m, k))\n# after\nsparse_group_elems = 8 if bitwidth(a_dtype) == 4 else 4\nmeta = build_metadata((m, k // sparse_group_elems * 2))","handlingStrategy":"validation","validationCode":"sge = 8 if utils.bitwidth(a_element_type) == 4 else 4\nassert a_sparse_metadata.shape == (m, k // sge * 2)","typeGuard":"def valid_sparse_meta(shape, m, k, dtype) -> bool:\n    sge = 8 if utils.bitwidth(dtype) == 4 else 4\n    return tuple(shape) == (m, k // sge * 2)","tryCatchPattern":null,"preventionTips":["Regenerate metadata whenever operand dtype bitwidth changes","Keep metadata shape derivation next to the pruning code"],"tags":["gpu","mosaic","tcgen05","sparsity","shape-mismatch","metadata"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}