{"record":{"id":"a60f28df6f4475e2","repo":"jax-ml/jax","slug":"side-contracting-dim-i-of-size-size-must-be","errorCode":null,"errorMessage":"{side} contracting dim {i} of size {size} must be divisible by its scale's dim size {scale.shape[i]}.","messagePattern":"(.+?) contracting dim (.+?) of size (.+?) must be divisible by its scale's dim size (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/scaled_dot.py","lineNumber":34,"sourceCode":"from collections.abc import Sequence\nimport jax\nfrom jax._src import core\nfrom jax._src import dispatch\nfrom jax._src import dtypes\nfrom jax._src import numpy as jnp\nfrom 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,","sourceCodeStart":16,"sourceCodeEnd":52,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/scaled_dot.py#L16-L52","documentation":"In jax._src.lax.scaled_dot, each operand's contracting dimension size must be an integer multiple of the corresponding dimension of its scale tensor. This validation runs in `_validate_operand_scale` before the scaled dot product is traced.","triggerScenarios":"Calling the scaled dot API with lhs contracting dim of size 6 and lhs_scale dim of size 4 (6 % 4 != 0); scales whose dims don't evenly tile the operand's contracting dims.","commonSituations":"Block-quantized matmuls where scale block size (e.g. 128) doesn't divide the hidden dimension (e.g. 1000); mismatched quantization granularity between operands and scale tensors.","solutions":["Reshape/pad the operand so each contracting dim is divisible by the scale dim","Regenerate the scale tensors with block sizes dividing the operand dims (e.g. choose per-tensor scales or smaller blocks)","Fix the model config so hidden dims are multiples of the quantization block size"],"exampleFix":"# before\nlhs.shape[-1] == 1000; lhs_scale.shape[-1] == 128\n# after\nlhs = pad_to(lhs, 1024)  # or use scale block 250/500","handlingStrategy":"validation","validationCode":"for i, s in enumerate(scale.shape):\n    if i in contracting_dims:\n        assert operand.shape[i] % s == 0, f'dim {i} not divisible'","typeGuard":"def contracting_dims_divisible(operand, scale, cdims) -> bool:\n    return all(operand.shape[i] % scale.shape[i] == 0 for i in cdims)","tryCatchPattern":null,"preventionTips":["Choose quantization block sizes that divide model hidden dims","Pad operands to multiples of block size"],"tags":["jax","scaled-dot","shape-validation","quantization"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}