jax-ml/jax · error · ValueError

A scale layout {a_scale.layout} is not supported

Error message

A scale layout {a_scale.layout} is not supported

What it means

The A scale tensor in tcgen05.mma must carry the specific TMEM layout returned by scales_layout(). Any other layout (or None) is rejected because the tensor-core instruction reads scales from fixed TMEM lanes.

Source

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

              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")
    if b_scale.shape[0] % 128 or b_scale.shape[0] < n * num_cta:
      raise ValueError(
          f"B scale shape[0] must be a multiple of 128 and >= N={n * num_cta},"
          f" got {b_scale.shape[0]}"
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Construct a_scale with the layout from tcgen05.scales_layout()
2. Check that a_scale is a TensorMemRefView-like object with .layout set, not a raw memref

Example fix

# before
a_scale = TensorMemRefView(get_buffer(), (TMEM_ROWS, k_scales), dtype, layout=None)
# after
a_scale = TensorMemRefView(get_buffer(), (TMEM_ROWS, k_scales), dtype, layout=scales_layout())

assert a_scale.layout == scales_layout(), 'bad A scale layout'

When it happens

Trigger: Passing a_scale whose .layout attribute is not tcgen05.scales_layout(), e.g. a plain (row-major) layout or a swizzled SMEM layout.

Common situations: Building scale memrefs with generic layouts or reusing operand layouts for scales; missing the layout wrapping step in a kernel authoring helper.

Related errors

• Expected B scales to have a M=64 collective layout, got {b_s
• Expected B scales to have a M=128 layout, got {b_scale.layou
• 4-bit block scaled MMA only supports K-fastest operands, but
• 4-bit block scaled MMA only supports K-fastest operands, but
• A address calculation for multiple M tiles

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