{"record":{"id":"b88138b6530ac1d6","repo":"jax-ml/jax","slug":"unsupported-input-dtype-ty","errorCode":null,"errorMessage":"Unsupported input dtype: {ty}","messagePattern":"Unsupported input dtype: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":95,"sourceCode":"  # Bit 6 is reserved\n  def get_input_encoding(ty):\n    if ty == f16:\n      assert acc_dtype in {f16, f32}\n      return 0\n    elif ty == ir.BF16Type.get():\n      assert acc_dtype == f32\n      return 1\n    elif ty == ir.Float8E4M3FNType.get():\n      assert acc_dtype in {f16, f32}\n      return 0\n    elif ty == ir.Float8E5M2Type.get():\n      assert acc_dtype in {f16, f32}\n      return 1\n    elif ty == ir.IntegerType.get_signless(8):  # Only s8 for now.\n      assert acc_dtype == i32\n      return 1\n    else:\n      raise NotImplementedError(f\"Unsupported input dtype: {ty}\")\n  a_type_val = get_input_encoding(a_dtype)\n  b_type_val = get_input_encoding(b_dtype)\n  desc |= (a_type_val << 7)   # A dtype, bits 7-9\n  desc |= (b_type_val << 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  # Bit 23 is reserved\n  if m % 16 or m > 256:\n    raise ValueError(f\"M must be a multiple of 16 and <= 256, got: {m}\")\n  desc |= (m >> 4) << 24  # M >> 4, bits 24-28\n  # Bit 29 is reserved\n  # We ignore max shift under .ws, bits 30-31\n  return arith.constant(ir.IntegerType.get_signless(32), desc)\n","sourceCodeStart":77,"sourceCodeEnd":113,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L77-L113","documentation":"In create_instr_descriptor's get_input_encoding, matrix input dtypes are limited to f16, tf32 (encoding 0 per assert context) and 8-bit signed integer (s8). Any other input dtype cannot be encoded into the instruction descriptor and raises NotImplementedError.","triggerScenarios":"Passing A or B matrices with dtype like bf16, s32, or f64 to a tcgen05 mma.","commonSituations":"Copying Hopper WGMMA configs that use bf16 inputs into tcgen05 paths without adjusting dtypes; defaulting arrays to bf16.","solutions":["Use f16 inputs (with f16/f32 accumulator) or s8 inputs (with i32 accumulator)","Convert operands before the MMA, not after","Check current JAX for newly supported input dtypes"],"exampleFix":"# before\na = memref.cast(..., ir.BFloat16Type.get())\n\n# after\na = memref.cast(..., ir.F16Type.get())","handlingStrategy":"type-guard","validationCode":"f16, tf32 = ir.F16Type.get(), ir.FloatTF32Type.get()\ns8 = ir.IntegerType.get_signless(8)\nassert a_dtype in (f16, tf32, s8) and b_dtype in (f16, tf32, s8)","typeGuard":"def is_supported_input_dtype(dt) -> bool:\n    return dt in (ir.F16Type.get(), ir.FloatTF32Type.get(), ir.IntegerType.get_signless(8))","tryCatchPattern":null,"preventionTips":["Cast operands to f16/s8 before entering tcgen05 paths","Keep dtype checks next to descriptor construction"],"tags":["tcgen05","matmul","mosaic","gpu","dtype"],"backgroundTag":"unsupported-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}