jax-ml/jax · error · NotImplementedError

Unsupported input dtype: {ty}

Error message

Unsupported input dtype: {ty}

What it means

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.

Source

Thrown at jax/experimental/mosaic/gpu/tcgen05.py:95

  # Bit 6 is reserved
  def get_input_encoding(ty):
    if ty == f16:
      assert acc_dtype in {f16, f32}
      return 0
    elif ty == ir.BF16Type.get():
      assert acc_dtype == f32
      return 1
    elif ty == ir.Float8E4M3FNType.get():
      assert acc_dtype in {f16, f32}
      return 0
    elif ty == ir.Float8E5M2Type.get():
      assert acc_dtype in {f16, f32}
      return 1
    elif ty == ir.IntegerType.get_signless(8):  # Only s8 for now.
      assert acc_dtype == i32
      return 1
    else:
      raise NotImplementedError(f"Unsupported input dtype: {ty}")
  a_type_val = get_input_encoding(a_dtype)
  b_type_val = get_input_encoding(b_dtype)
  desc |= (a_type_val << 7)   # A dtype, bits 7-9
  desc |= (b_type_val << 10)  # B dtype, bits 10-12
  # We ignore negate bits 13-14
  desc |= transpose_a << 15  # Transpose A
  desc |= transpose_b << 16  # Transpose B
  if n % 8 or n > 256:
    raise ValueError(f"N must be a multiple of 8 and <= 256, got: {n}")
  desc |= (n >> 3) << 17  # N, bits 17-22
  # Bit 23 is reserved
  if m % 16 or m > 256:
    raise ValueError(f"M must be a multiple of 16 and <= 256, got: {m}")
  desc |= (m >> 4) << 24  # M >> 4, bits 24-28
  # Bit 29 is reserved
  # We ignore max shift under .ws, bits 30-31
  return arith.constant(ir.IntegerType.get_signless(32), desc)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use f16 inputs (with f16/f32 accumulator) or s8 inputs (with i32 accumulator)
2. Convert operands before the MMA, not after
3. Check current JAX for newly supported input dtypes

Example fix

# before
a = memref.cast(..., ir.BFloat16Type.get())

# after
a = memref.cast(..., ir.F16Type.get())

f16, tf32 = ir.F16Type.get(), ir.FloatTF32Type.get()
s8 = ir.IntegerType.get_signless(8)
assert a_dtype in (f16, tf32, s8) and b_dtype in (f16, tf32, s8)

def is_supported_input_dtype(dt) -> bool:
    return dt in (ir.F16Type.get(), ir.FloatTF32Type.get(), ir.IntegerType.get_signless(8))

When it happens

Trigger: Passing A or B matrices with dtype like bf16, s32, or f64 to a tcgen05 mma.

Common situations: Copying Hopper WGMMA configs that use bf16 inputs into tcgen05 paths without adjusting dtypes; defaulting arrays to bf16.

Related errors

• Unsupported accumulator dtype: {acc_dtype}
• N must be a multiple of 8 and <= 256, got: {n}
• M must be a multiple of 16 and <= 256, got: {m}
• Sparse MMA unsupported for f32
• MMA with element type {elem_type_str} does not support block

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/b88138b6530ac1d6. Report an issue: GitHub.