{"record":{"id":"009ee6465bb6fb13","repo":"jax-ml/jax","slug":"sparse-mma-unsupported-for-f32","errorCode":null,"errorMessage":"Sparse MMA unsupported for f32","messagePattern":"Sparse MMA unsupported for f32","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":292,"sourceCode":"      # TODO: We only need to check this if N is the minormost dim in B.\n      if 8 * b_swizzle // utils.bitwidth(a_element_type) > n // n_lane_groups:\n        raise ValueError(\n            f\"Swizzle={b_swizzle} is too big for MMA with M=64. Try\"\n            \" lowering it.\"\n        )\n  else:\n    raise ValueError(f\"Only M=128 and M=64 are supported for MMA, but got M={m}\")\n  f32 = ir.F32Type.get()\n  f16 = ir.F16Type.get()\n  s32 = ir.IntegerType.get_signless(32)\n  elem_type_str = (\n      f\"{a_element_type}\"\n      if a_element_type == b_element_type\n      else f\"({a_element_type}, {b_element_type})\"\n  )\n  if a_element_type == f32 or a_element_type == ir.BF16Type.get():\n    if a_element_type == f32 and is_sparse:\n      raise NotImplementedError(\"Sparse MMA unsupported for f32\")\n    if is_scaled:\n      raise ValueError(\n          f\"MMA with element type {elem_type_str} does not support block scaling\"\n      )\n    if d.dtype != f32:\n      raise ValueError(\n          f\"MMA with element type {elem_type_str} only supports accumulators\"\n          f\" of type f32, but got: {d.dtype}\"\n      )\n  elif a_element_type == f16:\n    if is_scaled:\n      raise ValueError(\n          f\"MMA with element type {elem_type_str} does not support block scaling\"\n      )\n    if d.dtype != f16 and d.dtype != f32:\n      raise ValueError(\n          f\"MMA with element type {elem_type_str} only supports accumulators of\"\n          f\" type f32 or f16, but got: {d.dtype}\"","sourceCodeStart":274,"sourceCodeEnd":310,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L274-L310","documentation":"Sparse MMA (tcgen05 sparse tensor cores) is only defined for fp16 and fp8 input types; the hardware has no fp32 sparse instruction, so requesting is_sparse with f32 A operands raises NotImplementedError.","triggerScenarios":"Calling tcgen05.mma with lhs of dtype f32 while supplying sparse metadata / enabling the sparse path.","commonSituations":"Enabling sparsity on a mixed-precision pipeline still configured with f32 activations; migrating a sparse fp16 kernel to f32 for accumulation-accuracy experiments without disabling sparsity.","solutions":["Use fp16 or fp8 (Float8E5M2/E4M3FN) for the A operand when doing sparse MMA","Keep the A/B operands in f32 only for the dense path (disable sparse metadata)","If numeric range is the concern, use block-scaled fp8 instead of f32"],"exampleFix":"# before\nlhs_f32 = lhs.astype(jnp.float32)\ntcgen05.mma(lhs_f32, rhs, d, sparse_metadata=meta)  # raises\n# after\nlhs_f16 = lhs.astype(jnp.float16)\ntcgen05.mma(lhs_f16, rhs_f16, d, sparse_metadata=meta)","handlingStrategy":"validation","validationCode":"if is_sparse:\n    assert a_element_type not in (ir.F32Type.get(),), 'Sparse MMA unsupported for f32'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Gate sparsity on operand dtype: only f16/f8","Keep a config table mapping dtype -> allowed features (sparse, scaled)"],"tags":["jax","mosaic","gpu","sparse","dtype","mma","tcgen05"],"backgroundTag":"unsupported-dtype-for-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}