jax-ml/jax · error · TypeError

{side} contracting dim {i} of size {size} must be divisible

Error message

{side} contracting dim {i} of size {size} must be divisible by its scale's dim size {scale.shape[i]}.

What it means

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.

Source

Thrown at jax/_src/lax/scaled_dot.py:34

from collections.abc import Sequence
import jax
from jax._src import core
from jax._src import dispatch
from jax._src import dtypes
from jax._src import numpy as jnp
from jax._src.interpreters import batching
from jax._src.interpreters import mlir
from jax._src.lax import lax
from jax._src.typing import Array, DTypeLike


def _validate_operand_scale(
    side, operand, scale, contracting_dims: Sequence[int]
):
  for i, size in enumerate(operand.shape):
    if i in contracting_dims:
      if size % scale.shape[i] != 0:
        raise TypeError(
            f"{side} contracting dim {i} of size {size} must be divisible by "
            f"its scale's dim size {scale.shape[i]}."
        )
      s = size // scale.shape[i]
      if s < 2:
        raise TypeError(
            f"The ratio of {side} contracting dim {i} to its scale's dim size"
            f" ({s}) must be at least 2."
        )
    elif scale.shape[i] != size:
      raise TypeError(
          f"{side} dim {i} of size {size} does not match scale dim size "
          f"{scale.shape[i]}."
      )


def _scaled_dot_validate_inputs(
    lhs: Array,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape/pad the operand so each contracting dim is divisible by the scale dim
2. Regenerate the scale tensors with block sizes dividing the operand dims (e.g. choose per-tensor scales or smaller blocks)
3. Fix the model config so hidden dims are multiples of the quantization block size

Example fix

# before
lhs.shape[-1] == 1000; lhs_scale.shape[-1] == 128
# after
lhs = pad_to(lhs, 1024)  # or use scale block 250/500

for i, s in enumerate(scale.shape):
    if i in contracting_dims:
        assert operand.shape[i] % s == 0, f'dim {i} not divisible'

def contracting_dims_divisible(operand, scale, cdims) -> bool:
    return all(operand.shape[i] % scale.shape[i] == 0 for i in cdims)

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• The ratio of {side} contracting dim {i} to its scale's dim s
• {side} dim {i} of size {size} does not match scale dim size
• Unrecognized mode: {config.mode}.
• scan got `length` argument of {} which disagrees with leadin
• conv_general_dilated batch_group_count must divide lhs batch

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/a60f28df6f4475e2. Report an issue: GitHub.