jax-ml/jax · error · ValueError

scaled_matmul requires scales to match non-contract dimensio

Error message

scaled_matmul requires scales to match non-contract dimensions of inputs, but got shapes a: {a.shape}, b: {b.shape}, a_scales: {a_scales.shape}, b_scales: {b_scales.shape}

What it means

The final scaled_matmul shape check: a_scales' M dim (dim 1) must equal a's M dim count of blocks and b_scales' N dim must match b's — concretely M_as == M_a and N_bs == N_b per the checked dims. If the scale tensors' non-contract dims don't match the operands, this composite error is raised.

Source

Thrown at jax/_src/nn/functions.py:1378

    B_b, N_b, K_b = b.shape
    if K_a != K_b or B_a != B_b:
        raise ValueError(
            "scaled_matmul requires inputs a and b to have matching batch (B) "
            f"and contract (K) dimensions, but got shapes {a.shape} and "
            f"{b.shape}"
        )

    B_as, M_as, K_as = a_scales.shape
    B_bs, N_bs, K_bs = b_scales.shape
    if K_as != K_bs or B_as != B_bs:
        raise ValueError(
            "scaled_matmul requires scales to have matching batch (B) and "
            f"contract (K) dimensions, but got shapes {a_scales.shape} and "
            f"{b_scales.shape}"
        )

    if M_as != M_a or N_bs != N_b:
        raise ValueError(
            "scaled_matmul requires scales to match non-contract dimensions of "
            f"inputs, but got shapes a: {a.shape}, b: {b.shape}, a_scales: "
            f"{a_scales.shape}, b_scales: {b_scales.shape}"
        )

    preferred_element_type = dtypes.check_and_canonicalize_user_dtype(
        preferred_element_type, "scaled_matmul"
    )
    out = cudnn_scaled_matmul(
        a,
        b,
        a_scales,
        b_scales,
        preferred_element_type=preferred_element_type,
    )
    return out

def get_scaled_dot_general_config(mode: Literal['nvfp4', 'mxfp8'],

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Compute scales with dims (B, ceil(M/32), ceil(K/32)) and (B, ceil(N/32), ceil(K/32)) matching the operands' block counts
2. Regenerate both scales with the same quantizer/block-size utility used to quantize a and b

Example fix

// before
a = jnp.zeros((2, 128, 64))
a_s = jnp.ones((2, 8, 2), jnp.float8_e8m0fnu)  # wrong M blocks

// after
import math
a_s = jnp.ones((2, math.ceil(128/32), math.ceil(64/32)), jnp.float8_e8m0fnu)  # (2,4,2)

import math
BLOCK = 32
assert a_s.shape[1] == math.ceil(a.shape[1]/BLOCK)
assert b_s.shape[1] == math.ceil(b.shape[1]/BLOCK)

When it happens

Trigger: Passing scales of shape (B, 8, Kblocks) when the operand has M=128 with 32-wide blocks (4 blocks expected); mismatched block granularity between data and scales.

Common situations: Hand-building MX format scale tensors without the ceil-division helper; changing block_size in one place only; quantizing with a different block size than dequant/scaled-matmul assumes.

Related errors

• scaled_matmul requires scales to have matching batch (B) and
• scaled_matmul requires all inputs to be 3-dimensional arrays
• scaled_matmul requires inputs a and b to have matching batch
• {name} ndim should be {len(shape)}, but got {t.ndim}
• {name} shape should be {shape}: but got {t.shape}

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