{"record":{"id":"cda8c6a00cfef823","repo":"jax-ml/jax","slug":"side-dim-i-of-size-size-does-not-match-scale","errorCode":null,"errorMessage":"{side} dim {i} of size {size} does not match scale dim size {scale.shape[i]}.","messagePattern":"(.+?) dim (.+?) of size (.+?) does not match scale dim size (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/scaled_dot.py","lineNumber":45,"sourceCode":"\ndef _validate_operand_scale(\n    side, operand, scale, contracting_dims: Sequence[int]\n):\n  for i, size in enumerate(operand.shape):\n    if i in contracting_dims:\n      if size % scale.shape[i] != 0:\n        raise TypeError(\n            f\"{side} contracting dim {i} of size {size} must be divisible by \"\n            f\"its scale's dim size {scale.shape[i]}.\"\n        )\n      s = size // scale.shape[i]\n      if s < 2:\n        raise TypeError(\n            f\"The ratio of {side} contracting dim {i} to its scale's dim size\"\n            f\" ({s}) must be at least 2.\"\n        )\n    elif scale.shape[i] != size:\n      raise TypeError(\n          f\"{side} dim {i} of size {size} does not match scale dim size \"\n          f\"{scale.shape[i]}.\"\n      )\n\n\ndef _scaled_dot_validate_inputs(\n    lhs: Array,\n    rhs: Array,\n    lhs_scale: Array | None,\n    rhs_scale: Array | None,\n    *,\n    dimension_numbers: lax.DotDimensionNumbers,\n    preferred_element_type: DTypeLike | None,\n):\n  \"\"\"Validates the inputs to scaled_dot.\"\"\"\n  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = dimension_numbers\n\n  ndims = [lhs.ndim, rhs.ndim]","sourceCodeStart":27,"sourceCodeEnd":63,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/scaled_dot.py#L27-L63","documentation":"For non-contracting dimensions, scaled_dot requires the scale tensor's shape to exactly equal the operand's shape dim-by-dim (no broadcasting). `_validate_operand_scale` raises TypeError naming the mismatched dim.","triggerScenarios":"lhs has batch dim of size 8 but lhs_scale has batch dim 1; scales built with squeezed or broadcast batch dims.","commonSituations":"Reusing a scale tensor computed for a different batch size; squeezing singleton dims when saving/loading quantization scales.","solutions":["Broadcast the scale to the operand's non-contracting shape before the call (e.g. np.broadcast_to + reshape)","Regenerate scales with the correct batch dimension","Keep only contracting dims coarser; all others must match exactly"],"exampleFix":"# before\nlhs_scale = scale[0]  # missing batch dim\n# after\nlhs_scale = np.broadcast_to(scale, lhs.shape_pair)  # match non-contracting dims exactly\n# e.g. np.broadcast_to(scale[None], (B,) + scale.shape)","handlingStrategy":"type-guard","validationCode":"non_cd = [i for i in range(operand.ndim) if i not in contracting_dims]\nassert all(scale.shape[i] == operand.shape[i] for i in non_cd)","typeGuard":"def non_contracting_match(operand, scale, cdims) -> bool:\n    return all(scale.shape[i] == operand.shape[i]\n               for i in range(operand.ndim) if i not in cdims)","tryCatchPattern":null,"preventionTips":["np.broadcast_to scales to operand shape before the call","Keep batch dims in scales when saving checkpoints"],"tags":["jax","scaled-dot","shape-validation"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}