{"record":{"id":"f8d910c906455295","repo":"jax-ml/jax","slug":"scaled-matmul-requires-inputs-a-and-b-to-have-matc","errorCode":null,"errorMessage":"scaled_matmul requires inputs a and b to have matching batch (B) and contract (K) dimensions, but got shapes {a.shape} and {b.shape}","messagePattern":"scaled_matmul requires inputs a and b 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":1362,"sourceCode":"      Using fused cuDNN call on Blackwell GPUs:\n\n      >>> dtype = jnp.float8_e4m3fn\n      >>> a = jax.random.normal(jax.random.PRNGKey(1), (3, 128, 64), dtype=dtype)\n      >>> b = jax.random.normal(jax.random.PRNGKey(2), (3, 128, 64), dtype=dtype)\n      >>> a_scales = jnp.ones((3, 128, 4), dtype=jnp.float8_e8m0fnu)\n      >>> b_scales = jnp.ones((3, 128, 4), dtype=jnp.float8_e8m0fnu)\n      >>> scaled_matmul(a, b, a_scales, b_scales)  # doctest: +SKIP\n    \"\"\"\n    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: \"","sourceCodeStart":1344,"sourceCodeEnd":1380,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/nn/functions.py#L1344-L1380","documentation":"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.","triggerScenarios":"a.shape=(2,128,64), b.shape=(4,256,64) (batch mismatch) or a K of 64 vs b K of 128.","commonSituations":"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.","solutions":["Ensure lhs.shape[0] == rhs.shape[0] (or broadcast manually by tiling the smaller batch)","Ensure lhs.shape[2] == rhs.shape[2] (the shared K dimension)","If b is (B,K,N), transpose with b.swapaxes(1,2) before calling"],"exampleFix":"// before\na = jnp.zeros((2, 128, 64)); b = jnp.zeros((2, 64, 128))\njax.nn.scaled_matmul(a, b, a_s, b_s)  # K mismatch: 64 vs 128\n\n// after\nb = b.swapaxes(1, 2)  # (2, 128, 64)\njax.nn.scaled_matmul(a, b, a_s, b_s)","handlingStrategy":"validation","validationCode":"assert a.shape[0] == b.shape[0] and a.shape[2] == b.shape[2], (\n    f'B/K mismatch: {a.shape} vs {b.shape}')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep operands in (B,M,K)/(B,N,K) canonical layout; transpose at the boundary","Unit-test shapes of quantization pipelines"],"tags":["jax","nn","matmul","float8","shape-validation"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}