jax-ml/jax · error · TypeError

{side} dim {i} of size {size} does not match scale dim size

Error message

{side} dim {i} of size {size} does not match scale dim size {scale.shape[i]}.

What it means

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.

Source

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

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,
    rhs: Array,
    lhs_scale: Array | None,
    rhs_scale: Array | None,
    *,
    dimension_numbers: lax.DotDimensionNumbers,
    preferred_element_type: DTypeLike | None,
):
  """Validates the inputs to scaled_dot."""
  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = dimension_numbers

  ndims = [lhs.ndim, rhs.ndim]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Broadcast the scale to the operand's non-contracting shape before the call (e.g. np.broadcast_to + reshape)
2. Regenerate scales with the correct batch dimension
3. Keep only contracting dims coarser; all others must match exactly

Example fix

# before
lhs_scale = scale[0]  # missing batch dim
# after
lhs_scale = np.broadcast_to(scale, lhs.shape_pair)  # match non-contracting dims exactly
# e.g. np.broadcast_to(scale[None], (B,) + scale.shape)

non_cd = [i for i in range(operand.ndim) if i not in contracting_dims]
assert all(scale.shape[i] == operand.shape[i] for i in non_cd)

def non_contracting_match(operand, scale, cdims) -> bool:
    return all(scale.shape[i] == operand.shape[i]
               for i in range(operand.ndim) if i not in cdims)

When it happens

Trigger: lhs has batch dim of size 8 but lhs_scale has batch dim 1; scales built with squeezed or broadcast batch dims.

Common situations: Reusing a scale tensor computed for a different batch size; squeezing singleton dims when saving/loading quantization scales.

Related errors

• {side} contracting dim {i} of size {size} must be divisible
• The ratio of {side} contracting dim {i} to its scale's dim s
• scan got `length` argument of {} which disagrees with leadin
• conv_general_dilated batch_group_count must divide lhs batch
• conv_general_dilated rhs output feature dimension size must

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