{"record":{"id":"13efe5c7c7f1a56a","repo":"jax-ml/jax","slug":"expected-metadata-dtype-to-be-uint2-got-meta-dt","errorCode":null,"errorMessage":"Expected metadata dtype to be uint2, got: {meta.dtype}","messagePattern":"Expected metadata dtype to be uint2, got: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/helpers.py","lineNumber":201,"sourceCode":"  return decorator\n\n\ndef format_tcgen05_sparse_metadata(meta, operand_dtype):\n  \"\"\"Formats the sparse metadata for tcgen05.mma into the expected format.\n\n  See\n  https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-sparse-matrices-sparsity-selector-kind-f16-m128-256\n  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:","sourceCodeStart":183,"sourceCodeEnd":219,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/helpers.py#L183-L219","documentation":"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.","triggerScenarios":"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.","commonSituations":"Preparing sparse metadata outside JAX (numpy defaults to uint8/uint64) and forgetting the cast; porting from a different sparse format where metadata was wider.","solutions":["Cast before calling: meta = meta.astype(jax.numpy.uint2)","Generate metadata directly as uint2 (e.g. via arange % 4 cast to uint2 in tests)"],"exampleFix":"// before\nmeta = meta.astype(jnp.uint8)\nformatted = format_tcgen05_sparse_metadata(meta, jnp.float8_e4m3fn)\n\n// after\nmeta = meta.astype(jnp.uint2)\nformatted = format_tcgen05_sparse_metadata(meta, jnp.float8_e4m3fn)","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nif meta.dtype != jnp.uint2:\n    meta = meta.astype(jnp.uint2)","typeGuard":"def metadata_dtype_ok(meta) -> bool:\n    return str(meta.dtype) == 'uint2'","tryCatchPattern":"try:\n    return format_tcgen05_sparse_metadata(meta, dt)\nexcept ValueError:\n    return format_tcgen05_sparse_metadata(meta.astype(jnp.uint2), dt)","preventionTips":["Cast metadata to uint2 at load time","Write a small loader that asserts dtype and shape per the docstring"],"tags":["jax","pallas","mosaic-gpu","sparse","tcgen05","dtype","validation"],"backgroundTag":"wrong-dtype-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}