{"record":{"id":"0cfa97f073ebec0f","repo":"jax-ml/jax","slug":"scale-element-type-mismatch-expected-f8e8m0fnu-or","errorCode":null,"errorMessage":"Scale element type mismatch: expected f8e8m0fnu or f8e4m3fn, got {scale_element_type}","messagePattern":"Scale element type mismatch: expected f8e8m0fnu or f8e4m3fn, got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":457,"sourceCode":"            f\" {scale_block}\"\n        )\n    elif isinstance(a_element_type, ir.Float4E2M1FNType):\n      if isinstance(scale_element_type, ir.Float8E4M3FNType):\n        if base_scale_block != 16:\n          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(","sourceCodeStart":439,"sourceCodeEnd":475,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L439-L475","documentation":"tcgen05.mma block-scaled MMA only supports scale tensors of type f8e8m0fnu (MX-style exponent scales) or f8e4m3fn. This error fires when the scale operands have any other dtype.","triggerScenarios":"Passing a_scale/b_scale whose element type is e.g. Float32, Float8E4M3FN, or an integer type to mma() with block scaling enabled.","commonSituations":"Using fp32 scales from a reference implementation; dequantization code that upcast scales; scale tensors created with the wrong MLIR type when building a kernel.","solutions":["Cast scale tensors to f8e8m0fnu (preferred MX format) or f8e4m3fn before calling mma","Ensure the scale buffers are allocated with the matching dtype at kernel setup"],"exampleFix":"# before\na_scale = memref.cast(...)  # f32 scales\ntcgen05.mma(..., a_scale=a_scale, scale_block=32)\n# after\na_scale = memref.cast(..., f8e8m0fnu)\ntcgen05.mma(..., a_scale=a_scale, scale_block=32)","handlingStrategy":"validation","validationCode":"from jaxlib.mlir import ir\nassert a_scale.dtype in (ir.Float8E8M0FNUType.get(), ir.Float8E4M3FNType.get())","typeGuard":"def is_supported_scale_dtype(t) -> bool:\n    return isinstance(t, (ir.Float8E8M0FNUType, ir.Float8E4M3FNType))","tryCatchPattern":null,"preventionTips":["Centralize scale tensor creation with a fixed dtype helper","Fail fast on scale dtype at buffer allocation time"],"tags":["gpu","mosaic","tcgen05","dtype","block-scaling"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}