jax-ml/jax · error · NotImplementedError

Unsupported element type for block scaling: {a_element_type}

Error message

Unsupported element type for block scaling: {a_element_type}

What it means

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.

Source

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

          expected = 32 if is_sparse else 16
          raise ValueError(
              f"Scale block size mismatch: expected {expected}, got"
              f" {scale_block}"
          )
      elif isinstance(scale_element_type, ir.Float8E8M0FNUType):
        if base_scale_block not in (16, 32):
          expected = "32 or 64" if is_sparse else "16 or 32"
          raise ValueError(
              f"Scale block size mismatch: expected {expected}, got"
              f" {scale_block}"
          )
      else:
        raise ValueError(
            "Scale element type mismatch: expected f8e8m0fnu or f8e4m3fn, got"
            f" {scale_element_type}"
        )
    else:
      raise NotImplementedError(
          f"Unsupported element type for block scaling: {a_element_type}"
      )
    k_scales = k // scale_block
    if a_scale.shape != (TMEM_ROWS, k_scales):
      raise ValueError(
          f"A scale shape mismatch: expected ({TMEM_ROWS}, {k_scales}), got"
          f" {a_scale.shape}"
      )
    if a_scale.layout != scales_layout():
      raise ValueError(f"A scale layout {a_scale.layout} is not supported")
    if collective and m == 64:
      if b_scale.layout != b_scales_m64_collective_layout():
        raise ValueError(
            "Expected B scales to have a M=64 collective layout, got"
            f" {b_scale.layout}"
        )
    elif m == 128:
      if b_scale.layout != scales_layout():

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove a_scale/b_scale arguments if you don't need block scaling
2. Convert operands to a supported MX type (f8e4m3fn, f8e5m2, f8e8m0fnu-backed, or 4-bit float types) before the MMA

Example fix

# before
tcgen05.mma(a_bf16, b, d, a_scale=asc, b_scale=bsc, scale_block=32)
# after
tcgen05.mma(a_bf16, b, d)

supported = (ir.Float8E4M3FNType, ir.Float8E5M2Type, ...)  # f8/f4 types
assert a_scale is None or isinstance(a_element_type, supported)

def is_block_scalable(t) -> bool:
    import jax.experimental.mosaic.gpu.utils as u
    return u.bitwidth(t) in (4, 8) and 'Float' in type(t).__name__

try:
    tcgen05.mma(...)
except NotImplementedError:
    tcgen05.mma(a, b, d)  # fallback without scaling

When it happens

Trigger: Calling mma() with a_scale and b_scale where a's dtype is e.g. f16, bf16, or an integer type.

Common situations: Trying to attach MX scales to an FP16 GEMM; leftover scale arguments when downgrading a kernel from MXFP8 to BF16.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• A scale address calculation for multiple M tiles
• B scale address calculation for multiple N tiles
• MMA with element type {elem_type_str} does not support block
• Scale element type mismatch: expected f8e8m0fnu or f8e4m3fn,
• A scale shape mismatch: expected ({TMEM_ROWS}, {k_scales}),

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