{"record":{"id":"f7017c192ecb6aa4","repo":"jax-ml/jax","slug":"expected-the-trailing-dimension-of-the-metadata-to","errorCode":null,"errorMessage":"Expected the trailing dimension of the metadata to be 2, got: {meta.shape[-1]}","messagePattern":"Expected the trailing dimension of the metadata to be 2, got: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/helpers.py","lineNumber":209,"sourceCode":"  for the documentation of the required layouts. The array can be copied into\n  SMEM, from where ``plgpu.async_copy_sparse_metadata_to_tmem`` can be used to\n  copy it over to TMEM. The formatting of the array depends on the data type of\n  the operands to the sparse MMA operation.\n\n  Args:\n    meta: Metadata of shape (M, K // 4, 2).\n    dtype: Data type of MMA operands.\n  \"\"\"\n  if meta.dtype != dtypes.uint2:\n    raise ValueError(f\"Expected metadata dtype to be uint2, got: {meta.dtype}\")\n  if meta.ndim != 3:\n    raise ValueError(\n        \"Expected metadata to be 3-dimensional (M, K // 4, 2), but it is\"\n        f\" {meta.ndim}D\"\n    )\n  m, k, _2 = meta.shape\n  if _2 != 2:\n    raise ValueError(\n        \"Expected the trailing dimension of the metadata to be 2, got:\"\n        f\" {meta.shape[-1]}\"\n    )\n  k *= 2\n  bitsize = dtypes.itemsize_bits(operand_dtype)\n  if bitsize == 8:\n    meta_tiled = meta.reshape(m // 128, 128, k // 64, 64).transpose(0, 2, 1, 3)\n  elif bitsize == 16:\n    meta_tiled = meta.reshape(m // 128, 8, 2, 8, k // 64, 4, 2, 8).transpose(0, 4, 1, 6, 3, 5, 2, 7)\n  else:\n    raise NotImplementedError(\n        f\"Sparse metadata format not implemented for {operand_dtype=}\"\n    )\n  return meta_tiled.reshape(m // 128, k // 64, 128, 64)\n\n\ndef find_swizzle(minor_dim_bits: int, what: str = \"\"):\n  \"\"\"Returns the largest swizzle that can be applied to a memory region.","sourceCodeStart":191,"sourceCodeEnd":227,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/helpers.py#L191-L227","documentation":"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.","triggerScenarios":"meta.shape[-1] != 2, e.g. reshaping to (M, K//4, 4) or (M, K//2, 1) before calling format_tcgen05_sparse_metadata.","commonSituations":"Mistaking the layout as (M, K//4, 1) pairs or interleaving metadata differently when preprocessing weights for sparse MMA.","solutions":["Fix the packing so the final axis holds the 2 metadata entries: meta.reshape(m, k // 4, 2)","Double-check against the docstring: shape (M, K // 4, 2), dtype uint2"],"exampleFix":"// before\nmeta = meta.reshape(m, k // 2, 1)\n\n// after\nmeta = meta.reshape(m, k // 4, 2)","handlingStrategy":"validation","validationCode":"assert meta.shape[-1] == 2, f'trailing dim must be 2, got {meta.shape[-1]}'","typeGuard":"def trailing_dim_is_2(meta) -> bool:\n    return meta.ndim >= 1 and meta.shape[-1] == 2","tryCatchPattern":"try:\n    return format_tcgen05_sparse_metadata(meta, dt)\nexcept ValueError:\n    meta = meta.reshape(*meta.shape[:-1], -1)\n    meta = meta.reshape(*meta.shape[:-1], meta.shape[-1] // 2, 2)\n    return format_tcgen05_sparse_metadata(meta, dt)","preventionTips":["Follow the documented (M, K // 4, 2) packing exactly","Validate with a shape assert in test fixtures"],"tags":["jax","pallas","mosaic-gpu","sparse","tcgen05","shape","last-dim"],"backgroundTag":"wrong-array-shape","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}