jax-ml/jax · error · ValueError

scaled_matmul requires scales to have matching batch (B) and

Error message

scaled_matmul requires scales to have matching batch (B) and contract (K) dimensions, but got shapes {a_scales.shape} and {b_scales.shape}

What it means

scaled_matmul requires the two scale tensors to agree on their batch (dim 0) and contraction (dim 2) dimensions, mirroring the operand constraint. If a_scales.shape[0/2] != b_scales.shape[0/2] this is raised.

Source

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

    a, b, a_scales, b_scales = lhs, rhs, lhs_scales, rhs_scales
    if not all(x.ndim == 3 for x in (a, b, a_scales, b_scales)):
        raise ValueError(
            "scaled_matmul requires all inputs to be 3-dimensional arrays"
        )

    B_a, M_a, K_a = a.shape
    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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make both scale tensors' B and K dims identical, typically (B, ceil(M or N / block), ceil(K / block))
Use a shared helper to allocate both scale tensors with the same block size

Example fix

// before
a_s = jnp.ones((1, 4, 2), jnp.float8_e8m0fnu)
b_s = jnp.ones((2, 4, 2), jnp.float8_e8m0fnu)

// after
B = 2
a_s = jnp.ones((B, 4, 2), jnp.float8_e8m0fnu)
b_s = jnp.ones((B, 4, 2), jnp.float8_e8m0fnu)

Defensive patterns

Strategy: validation

Validate before calling

assert a_s.shape[0] == b_s.shape[0] and a_s.shape[2] == b_s.shape[2], (
    f'scale B/K mismatch: {a_s.shape} vs {b_s.shape}')

Prevention

Allocate both scale tensors in one place with a shared block size
Store block_size alongside tensors to keep scale shapes consistent

When it happens

Trigger: a_scales=(1,4,2) and b_scales=(2,4,2) (batch mismatch), or scale block counts along K differing because operands used different block sizes.

Common situations: Generating scales with different granularity (32 vs 128 blocks) for lhs and rhs; broadcasting one scale tensor across batches but not the other; unit test fixtures with inconsistent random shapes.

Related errors

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