{"record":{"id":"9906210ae55f2bf2","repo":"jax-ml/jax","slug":"the-ratio-of-side-contracting-dim-i-to-its-sca","errorCode":null,"errorMessage":"The ratio of {side} contracting dim {i} to its scale's dim size ({s}) must be at least 2.","messagePattern":"The ratio of (.+?) contracting dim (.+?) to its scale's dim size \\((.+?)\\) must be at least 2\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/scaled_dot.py","lineNumber":40,"sourceCode":"from jax._src.interpreters import batching\nfrom jax._src.interpreters import mlir\nfrom jax._src.lax import lax\nfrom jax._src.typing import Array, DTypeLike\n\n\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,","sourceCodeStart":22,"sourceCodeEnd":58,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/scaled_dot.py#L22-L58","documentation":"Even when divisible, scaled_dot requires the ratio between an operand's contracting dim and its scale's dim to be at least 2 — i.e. the scale must actually compress the dimension. A ratio of 1 (same size) is rejected as it makes scaling meaningless.","triggerScenarios":"Passing a scale tensor whose contracting dim equals the operand's contracting dim (ratio 1), e.g. a full-size per-element scale on a contracting axis.","commonSituations":"Migrating from per-element scaling code to the block-scaled API; accidentally broadcasting a scale to full operand shape.","solutions":["Use a coarser scale (block size such that ratio >= 2) on contracting dims","If you need per-element scaling on contracting dims, use standard dot with explicit multiplication instead","Double-check which dims are marked contracting in the dimension_numbers"],"exampleFix":"# before\nscale.shape[k] == operand.shape[k]  # ratio 1\n# after\nscale = reshape_block_scale(scale, block=2)  # ratio >= 2","handlingStrategy":"validation","validationCode":"assert all(operand.shape[i] // scale.shape[i] >= 2 for i in contracting_dims)","typeGuard":"def ratios_at_least_2(operand, scale, cdims) -> bool:\n    return all(operand.shape[i] // scale.shape[i] >= 2 for i in cdims)","tryCatchPattern":null,"preventionTips":["Never use full-size scales on contracting dims; use block scales"],"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"}