{"record":{"id":"b0f07ea203121eb9","repo":"jax-ml/jax","slug":"unsupported-element-type-for-block-scaling-a-ele","errorCode":null,"errorMessage":"Unsupported element type for block scaling: {a_element_type}","messagePattern":"Unsupported element type for block scaling: (.+?)","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":462,"sourceCode":"          expected = 32 if is_sparse else 16\n          raise ValueError(\n              f\"Scale block size mismatch: expected {expected}, got\"\n              f\" {scale_block}\"\n          )\n      elif isinstance(scale_element_type, ir.Float8E8M0FNUType):\n        if base_scale_block not in (16, 32):\n          expected = \"32 or 64\" if is_sparse else \"16 or 32\"\n          raise ValueError(\n              f\"Scale block size mismatch: expected {expected}, got\"\n              f\" {scale_block}\"\n          )\n      else:\n        raise ValueError(\n            \"Scale element type mismatch: expected f8e8m0fnu or f8e4m3fn, got\"\n            f\" {scale_element_type}\"\n        )\n    else:\n      raise NotImplementedError(\n          f\"Unsupported element type for block scaling: {a_element_type}\"\n      )\n    k_scales = k // scale_block\n    if a_scale.shape != (TMEM_ROWS, k_scales):\n      raise ValueError(\n          f\"A scale shape mismatch: expected ({TMEM_ROWS}, {k_scales}), got\"\n          f\" {a_scale.shape}\"\n      )\n    if a_scale.layout != scales_layout():\n      raise ValueError(f\"A scale layout {a_scale.layout} is not supported\")\n    if collective and m == 64:\n      if b_scale.layout != b_scales_m64_collective_layout():\n        raise ValueError(\n            \"Expected B scales to have a M=64 collective layout, got\"\n            f\" {b_scale.layout}\"\n        )\n    elif m == 128:\n      if b_scale.layout != scales_layout():","sourceCodeStart":444,"sourceCodeEnd":480,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L444-L480","documentation":"Block scaling in tcgen05.mma is only implemented for 4-bit (f4/nf4) and 8-bit (f8) floating point A operand types. Any other A element type with a_scale/b_scale supplied raises NotImplementedError.","triggerScenarios":"Calling mma() with a_scale and b_scale where a's dtype is e.g. f16, bf16, or an integer type.","commonSituations":"Trying to attach MX scales to an FP16 GEMM; leftover scale arguments when downgrading a kernel from MXFP8 to BF16.","solutions":["Remove a_scale/b_scale arguments if you don't need block scaling","Convert operands to a supported MX type (f8e4m3fn, f8e5m2, f8e8m0fnu-backed, or 4-bit float types) before the MMA"],"exampleFix":"# before\ntcgen05.mma(a_bf16, b, d, a_scale=asc, b_scale=bsc, scale_block=32)\n# after\ntcgen05.mma(a_bf16, b, d)","handlingStrategy":"validation","validationCode":"supported = (ir.Float8E4M3FNType, ir.Float8E5M2Type, ...)  # f8/f4 types\nassert a_scale is None or isinstance(a_element_type, supported)","typeGuard":"def is_block_scalable(t) -> bool:\n    import jax.experimental.mosaic.gpu.utils as u\n    return u.bitwidth(t) in (4, 8) and 'Float' in type(t).__name__","tryCatchPattern":"try:\n    tcgen05.mma(...)\nexcept NotImplementedError:\n    tcgen05.mma(a, b, d)  # fallback without scaling","preventionTips":["Only pass scales for MX-type operands","Gate scale arguments behind an is_mx(dtype) check"],"tags":["gpu","mosaic","tcgen05","block-scaling","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}