jax-ml/jax · error · TypeError

The ratio of {side} contracting dim {i} to its scale's dim s

Error message

The ratio of {side} contracting dim {i} to its scale's dim size ({s}) must be at least 2.

What it means

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.

Source

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

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,
    rhs: Array,
    lhs_scale: Array | None,
    rhs_scale: Array | None,
    *,
    dimension_numbers: lax.DotDimensionNumbers,
    preferred_element_type: DTypeLike | None,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a coarser scale (block size such that ratio >= 2) on contracting dims
2. If you need per-element scaling on contracting dims, use standard dot with explicit multiplication instead
3. Double-check which dims are marked contracting in the dimension_numbers

Example fix

# before
scale.shape[k] == operand.shape[k]  # ratio 1
# after
scale = reshape_block_scale(scale, block=2)  # ratio >= 2

assert all(operand.shape[i] // scale.shape[i] >= 2 for i in contracting_dims)

def ratios_at_least_2(operand, scale, cdims) -> bool:
    return all(operand.shape[i] // scale.shape[i] >= 2 for i in cdims)

When it happens

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

Common situations: Migrating from per-element scaling code to the block-scaled API; accidentally broadcasting a scale to full operand shape.

Related errors

• {side} contracting dim {i} of size {size} must be divisible
• {side} dim {i} of size {size} does not match scale dim size
• 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/9906210ae55f2bf2. Report an issue: GitHub.