{"record":{"id":"5c8d0cf1eab21830","repo":"jax-ml/jax","slug":"scaled-matmul-requires-scales-to-match-non-contrac","errorCode":null,"errorMessage":"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}","messagePattern":"scaled_matmul requires scales to match non-contract dimensions of inputs, but got shapes a: (.+?), b: (.+?), a_scales: (.+?), b_scales: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/nn/functions.py","lineNumber":1378,"sourceCode":"    B_b, N_b, K_b = b.shape\n    if K_a != K_b or B_a != B_b:\n        raise ValueError(\n            \"scaled_matmul requires inputs a and b to have matching batch (B) \"\n            f\"and contract (K) dimensions, but got shapes {a.shape} and \"\n            f\"{b.shape}\"\n        )\n\n    B_as, M_as, K_as = a_scales.shape\n    B_bs, N_bs, K_bs = b_scales.shape\n    if K_as != K_bs or B_as != B_bs:\n        raise ValueError(\n            \"scaled_matmul requires scales to have matching batch (B) and \"\n            f\"contract (K) dimensions, but got shapes {a_scales.shape} and \"\n            f\"{b_scales.shape}\"\n        )\n\n    if M_as != M_a or N_bs != N_b:\n        raise ValueError(\n            \"scaled_matmul requires scales to match non-contract dimensions of \"\n            f\"inputs, but got shapes a: {a.shape}, b: {b.shape}, a_scales: \"\n            f\"{a_scales.shape}, b_scales: {b_scales.shape}\"\n        )\n\n    preferred_element_type = dtypes.check_and_canonicalize_user_dtype(\n        preferred_element_type, \"scaled_matmul\"\n    )\n    out = cudnn_scaled_matmul(\n        a,\n        b,\n        a_scales,\n        b_scales,\n        preferred_element_type=preferred_element_type,\n    )\n    return out\n\ndef get_scaled_dot_general_config(mode: Literal['nvfp4', 'mxfp8'],","sourceCodeStart":1360,"sourceCodeEnd":1396,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/nn/functions.py#L1360-L1396","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Compute scales with dims (B, ceil(M/32), ceil(K/32)) and (B, ceil(N/32), ceil(K/32)) matching the operands' block counts","Regenerate both scales with the same quantizer/block-size utility used to quantize a and b"],"exampleFix":"// before\na = jnp.zeros((2, 128, 64))\na_s = jnp.ones((2, 8, 2), jnp.float8_e8m0fnu)  # wrong M blocks\n\n// after\nimport math\na_s = jnp.ones((2, math.ceil(128/32), math.ceil(64/32)), jnp.float8_e8m0fnu)  # (2,4,2)","handlingStrategy":"validation","validationCode":"import math\nBLOCK = 32\nassert a_s.shape[1] == math.ceil(a.shape[1]/BLOCK)\nassert b_s.shape[1] == math.ceil(b.shape[1]/BLOCK)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use one quantization utility to emit both operands and scales together","Derive scale shapes via ceil-division helpers, never hardcode them"],"tags":["jax","nn","matmul","float8","scales","shape-validation"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}