jax-ml/jax · error · ValueError

scaled_matmul requires inputs a and b to have matching batch

Error message

scaled_matmul requires inputs a and b to have matching batch (B) and contract (K) dimensions, but got shapes {a.shape} and {b.shape}

What it means

In scaled_matmul the two operands must agree on the batch dimension B (dim 0) and the contraction dimension K (dim 2, for both (B,M,K) and (B,N,K)). Mismatch raises this with both shapes printed.

Source

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

      Using fused cuDNN call on Blackwell GPUs:

      >>> dtype = jnp.float8_e4m3fn
      >>> a = jax.random.normal(jax.random.PRNGKey(1), (3, 128, 64), dtype=dtype)
      >>> b = jax.random.normal(jax.random.PRNGKey(2), (3, 128, 64), dtype=dtype)
      >>> a_scales = jnp.ones((3, 128, 4), dtype=jnp.float8_e8m0fnu)
      >>> b_scales = jnp.ones((3, 128, 4), dtype=jnp.float8_e8m0fnu)
      >>> scaled_matmul(a, b, a_scales, b_scales)  # doctest: +SKIP
    """
    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: "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure lhs.shape[0] == rhs.shape[0] (or broadcast manually by tiling the smaller batch)
2. Ensure lhs.shape[2] == rhs.shape[2] (the shared K dimension)
3. If b is (B,K,N), transpose with b.swapaxes(1,2) before calling

Example fix

// before
a = jnp.zeros((2, 128, 64)); b = jnp.zeros((2, 64, 128))
jax.nn.scaled_matmul(a, b, a_s, b_s)  # K mismatch: 64 vs 128

// after
b = b.swapaxes(1, 2)  # (2, 128, 64)
jax.nn.scaled_matmul(a, b, a_s, b_s)

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

When it happens

Trigger: a.shape=(2,128,64), b.shape=(4,256,64) (batch mismatch) or a K of 64 vs b K of 128.

Common situations: Reusing matrices from a different microbatch size; transposing b incorrectly so N and K are swapped; wrong block granularity producing off-by-one K in scales-versus-operand checks upstream.

Related errors

• scaled_matmul requires all inputs to be 3-dimensional arrays
• scaled_matmul requires scales to have matching batch (B) and
• scaled_matmul requires scales to match non-contract dimensio
• {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/f8d910c906455295. Report an issue: GitHub.