jax-ml/jax · error · ValueError

B scale shape mismatch: expected ({b_scale.shape[0]}, {k_sca

Error message

B scale shape mismatch: expected ({b_scale.shape[0]}, {k_scales}), got {b_scale.shape}

What it means

In block-scaled MMA, the B scale tensor's second dimension must equal k_scales = k // scale_block. This fires when b_scale.shape[1] doesn't match, i.e. the scale count along K differs from what the operands imply.

Source

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

      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]}"
      )
    if b_scale.shape[1] != k_scales:
      raise ValueError(
          f"B scale shape mismatch: expected ({b_scale.shape[0]}, {k_scales}),"
          f" got {b_scale.shape}"
      )
  if is_sparse:
    sparse_group_elems = 8 if utils.bitwidth(a_element_type) == 4 else 4
    # Each sparse group has 2 entries.
    expected_meta_k = k // sparse_group_elems * 2
    if a_sparse_metadata.shape != (m, expected_meta_k):
      raise ValueError(
          f"A sparse metadata shape mismatch: expected {(m, expected_meta_k)},"
          f" got {a_sparse_metadata.shape}"
      )
    if a_sparse_metadata.dtype != ir.IntegerType.get_signless(2):
      raise ValueError(
          "A sparse metadata dtype mismatch: expected i2, got"
          f" {a_sparse_metadata.dtype}"
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Regenerate b_scale so shape[1] == k // scale_block
2. Verify B operand's K dimension and that A and B scales use the same scale_block

Example fix

# before
b_scale = make_scales((512, 8))
# after
k_scales = k // scale_block
b_scale = make_scales((512, k_scales))

assert b_scale.shape[1] == k // scale_block

def valid_b_scale(shape, k, scale_block) -> bool:
    return shape[1] == k // scale_block

When it happens

Trigger: b_scale.shape[1] != k // scale_block, e.g. k=512, scale_block=32 (k_scales=16) but b_scale has shape (..., 8).

Common situations: Transposed B operands where K and N are swapped in the scale tensor; inconsistent scale_block between the A and B scale computation.

Related errors

• A scale shape mismatch: expected ({TMEM_ROWS}, {k_scales}),
• B scale shape[0] must be a multiple of 128 and >= N={n * num
• 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/c6531d92eb0551dc. Report an issue: GitHub.