jax-ml/jax · error · ValueError

A scale shape mismatch: expected ({TMEM_ROWS}, {k_scales}),

Error message

A scale shape mismatch: expected ({TMEM_ROWS}, {k_scales}), got {a_scale.shape}

What it means

In block-scaled MMA, the A scale tensor must have exactly shape (TMEM_ROWS, k // scale_block). This error fires when a_scale.shape deviates, i.e. the scale tile doesn't cover K exactly once per scale block.

Source

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

      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():
        raise ValueError(
            f"Expected B scales to have a M=128 layout, got {b_scale.layout}"
        )
    else:
      raise AssertionError("Should not happen")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Regenerate a_scale so its shape is (TMEM_ROWS, k // scale_block)
2. Ensure k is a multiple of scale_block and that the same k is used for operands and scales

Example fix

# before
k, sb = 256, 32
a_scale = make_scale((TMEM_ROWS, 4))
# after
a_scale = make_scale((TMEM_ROWS, k // sb))

TMEM_ROWS = 128
assert a_scale.shape == (TMEM_ROWS, k // scale_block)

def valid_a_scale(shape, k, scale_block, rows=128) -> bool:
    return tuple(shape) == (rows, k // scale_block)

When it happens

Trigger: Passing a_scale with shape != (TMEM_ROWS, k/scale_block), e.g. k=256, scale_block=32 but a_scale shape (128, 4) instead of (128, 8).

Common situations: Scale tensors computed for a different K dimension than the operands; off-by-one in k_scales computation in a code generator; padded K without padding the scales.

Related errors

• B scale shape[0] must be a multiple of 128 and >= N={n * num
• B scale shape mismatch: expected ({b_scale.shape[0]}, {k_sca
• MMA with element type {elem_type_str} does not support block
• Scale element type mismatch: expected f8e8m0fnu or f8e4m3fn,
• Unsupported element type for block scaling: {a_element_type}

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