{"record":{"id":"5e601cbc1c5062d8","repo":"jax-ml/jax","slug":"unsupported-scale-type-scale-type","errorCode":null,"errorMessage":"Unsupported scale type: {scale_type}","messagePattern":"Unsupported scale type: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":148,"sourceCode":"  desc |= sparse << 2  # Sparsity, bit 2\n  # Bit 3 is reserved\n  assert 0 <= b_scale_idx < 4\n  desc |= b_scale_idx << 4  # B scale factor data ID, bits 4-5\n  # Bit 6 is reserved\n  desc |= get_input_encoding(a_type) << 7  # A dtype, bits 7-9\n  desc |= get_input_encoding(b_type) << 10  # B dtype, bits 10-12\n  # We ignore negate bits 13-14\n  desc |= transpose_a << 15  # Transpose A\n  desc |= transpose_b << 16  # Transpose B\n  if n % 8 or n > 256:\n    raise ValueError(f\"N must be a multiple of 8 and <= 256, got: {n}\")\n  desc |= (n >> 3) << 17  # N, bits 17-22\n  if scale_type == ir.Float8E8M0FNUType.get():\n    scale_encoding = 1\n  elif scale_type == ir.Float8E4M3FNType.get():\n    scale_encoding = 0\n  else:\n    raise NotImplementedError(f\"Unsupported scale type: {scale_type}\")\n  desc |= scale_encoding << 23  # Scale matrix type\n  # Bits 24-26 are reserved\n  if m % 128 or m > 256:\n    raise ValueError(f\"M must be a multiple of 16 and <= 256, got: {m}\")\n  desc |= (m >> 7) << 27  # M >> 7, bits 27-28\n  desc |= a_scale_idx << 29  # A scale factor data ID, bits 29-30\n  # Bit 31 is reserved\n  return arith.constant(ir.IntegerType.get_signless(32), desc)\n\n\ndef create_scaled_f8f6f4_instr_descriptor(*args, **kwargs) -> ir.Value:\n  def get_input_encoding(ty):\n    if ty == ir.Float8E4M3FNType.get():\n      return 0\n    elif ty == ir.Float8E5M2Type.get():\n      return 1\n    else:\n      raise NotImplementedError(f\"Unsupported input dtype: {ty}\")","sourceCodeStart":130,"sourceCodeEnd":166,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L130-L166","documentation":"For block-scaled MMA, the scale matrix dtype is encoded in descriptor bit 23: only Float8E8M0 (encoding 1, the MX standard) and Float8E4M3FN (encoding 0) are supported hardware scale formats. Other scale dtypes raise NotImplementedError.","triggerScenarios":"Passing scale_type other than f8e8m0 (e4m3 is accepted) — e.g. f32 scales or e5m2 scales — to the scaled descriptor builders.","commonSituations":"Writing custom scaled matmuls with FP32 scales instead of the MX-standard UE8M0; mixing up scale element formats.","solutions":["Use Float8E8M0FNUType scales (the MX standard)","Or Float8E4M3FNType if that variant is intended","Quantize/convert f32 scales to e8m0 before the MMA"],"exampleFix":"# before\nscale = tmem.alloc(dtype=ir.F32Type.get(), ...)\n\n# after\nscale = tmem.alloc(dtype=ir.Float8E8M0FNUType.get(), ...)","handlingStrategy":"type-guard","validationCode":"ok = (ir.Float8E8M0FNUType.get(), ir.Float8E4M3FNType.get())\nassert scale_type in ok, f'unsupported scale type {scale_type}'","typeGuard":"def is_supported_scale_type(dt) -> bool:\n    return dt in (ir.Float8E8M0FNUType.get(), ir.Float8E4M3FNType.get())","tryCatchPattern":null,"preventionTips":["Standardize on UE8M0 scales (MX standard)","Convert f32 scales to e8m0 during quantization"],"tags":["tcgen05","mx-format","matmul","dtype"],"backgroundTag":"unsupported-scale-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}