jax-ml/jax · error · TypeError

LHS and RHS must have the same number of batch dimensions, g

Error message

LHS and RHS must have the same number of batch dimensions, got {len(lhs_batch)} and {len(rhs_batch)}.

What it means

scaled_dot's dimension_numbers must declare the same number of batch dimensions for lhs and rhs. When `len(lhs_batch) != len(rhs_batch)` the validation raises TypeError with both counts.

Source

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

  """Validates the inputs to scaled_dot."""
  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = dimension_numbers

  ndims = [lhs.ndim, rhs.ndim]
  if lhs_scale is not None:
    ndims.append(lhs_scale.ndim)
  if rhs_scale is not None:
    ndims.append(rhs_scale.ndim)

  if max(ndims) != min(ndims):
    raise TypeError(
        "All input tensors must have the same rank. Got lhs rank:"
        f" {lhs.ndim} rhs rank: {rhs.ndim} lhs_scale rank:"
        f" {lhs_scale.ndim if lhs_scale is not None else 'N/A'} rhs_scale"
        f" rank: {rhs_scale.ndim if rhs_scale is not None else 'N/A'}."
    )

  if len(lhs_batch) != len(rhs_batch):
    raise TypeError(
        "LHS and RHS must have the same number of batch dimensions, got"
        f" {len(lhs_batch)} and {len(rhs_batch)}."
    )
  if len(lhs_contracting) != len(rhs_contracting):
    raise TypeError(
        "LHS and RHS must have the same number of contracting dimensions, got"
        f" {len(lhs_contracting)} and {len(rhs_contracting)}."
    )

  for i_lhs, i_rhs in zip(lhs_batch, rhs_batch):
    batch_dims_sizes = [
        lhs.shape[i_lhs],
        rhs.shape[i_rhs],
    ]
    if lhs_scale is not None:
      batch_dims_sizes.append(lhs_scale.shape[i_lhs])
    if rhs_scale is not None:
      batch_dims_sizes.append(rhs_scale.shape[i_rhs])

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the batch dimension lists equal length, typically both empty for plain matmuls or both containing the batch axis
2. Prefer passing dimension numbers as the standard ((lc, rc), (lb, rb)) tuple format and double-check symmetry
3. Test with the same shapes you use in an equivalent jnp.einsum to derive correct dims

Example fix

# before
dims = ((2,), (1,), (0,), ())  # lhs has batch, rhs doesn't
# after
dims = ((2,), (1,), (0,), (0,))  # matching batch dims

assert len(lhs_batch) == len(rhs_batch), 'batch dim counts must match'
assert len(lhs_contract) == len(rhs_contract)

def valid_dim_numbers(dims) -> bool:
    (lc, rc), (lb, rb) = dims
    return len(lb) == len(rb) and len(lc) == len(rc)

When it happens

Trigger: Passing dimension numbers where lhs has 1 batch dim and rhs has 0, e.g. ((0,), (), (2,), (1,)) with mismatched batch lists; forgetting the rhs batch dim in a custom dimension_numbers tuple.

Common situations: Hand-constructing lax-style dimension numbers ((lhs_contract, rhs_contract), (lhs_batch, rhs_batch)); adapting dot_general dimension numbers to the scaled variant and dropping one entry.

Related errors

• No 4+ dimensional dimension_number defaults.
• convolution dimension_numbers list/tuple must be length 3, g
• convolution dimension_numbers elements must be strings, got
• convolution dimension_numbers[{}] must have len equal to the
• convolution dimension_numbers must be tuple/list or None, go

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