{"record":{"id":"db4938b4f5683f5c","repo":"jax-ml/jax","slug":"scaled-matmul-requires-scales-to-have-matching-bat","errorCode":null,"errorMessage":"scaled_matmul requires scales to have matching batch (B) and contract (K) dimensions, but got shapes {a_scales.shape} and {b_scales.shape}","messagePattern":"scaled_matmul requires scales to have matching batch \\(B\\) and contract \\(K\\) dimensions, but got shapes (.+?) and (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/nn/functions.py","lineNumber":1371,"sourceCode":"    a, b, a_scales, b_scales = lhs, rhs, lhs_scales, rhs_scales\n    if not all(x.ndim == 3 for x in (a, b, a_scales, b_scales)):\n        raise ValueError(\n            \"scaled_matmul requires all inputs to be 3-dimensional arrays\"\n        )\n\n    B_a, M_a, K_a = a.shape\n    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,","sourceCodeStart":1353,"sourceCodeEnd":1389,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/nn/functions.py#L1353-L1389","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"// before\na_s = jnp.ones((1, 4, 2), jnp.float8_e8m0fnu)\nb_s = jnp.ones((2, 4, 2), jnp.float8_e8m0fnu)\n\n// after\nB = 2\na_s = jnp.ones((B, 4, 2), jnp.float8_e8m0fnu)\nb_s = jnp.ones((B, 4, 2), jnp.float8_e8m0fnu)","handlingStrategy":"validation","validationCode":"assert a_s.shape[0] == b_s.shape[0] and a_s.shape[2] == b_s.shape[2], (\n    f'scale B/K mismatch: {a_s.shape} vs {b_s.shape}')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Allocate both scale tensors in one place with a shared block size","Store block_size alongside tensors to keep scale shapes consistent"],"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"}