jax-ml/jax · error · TypeError

All input tensors must have the same rank. Got lhs rank: {lh

Error message

All input tensors must have the same rank. Got lhs rank: {lhs.ndim} rhs rank: {rhs.ndim} lhs_scale rank: {lhs_scale.ndim if lhs_scale is not None else 'N/A'} rhs_scale rank: {rhs_scale.ndim if rhs_scale is not None else 'N/A'}.

What it means

scaled_dot requires lhs, rhs and any provided lhs_scale/rhs_scale to all have the same rank (number of dimensions). `_scaled_dot_validate_inputs` compares min/max of the collected ndims and raises TypeError listing each rank when they differ.

Source

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

    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]
  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):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Expand dims so every input has the same rank (e.g. scale[None] for batched operands)
2. Reshape/unsqueeze operands or scales consistently
3. Verify all tensors' .ndim right before the call in tests

Example fix

# before
lhs.ndim == 3, lhs_scale.ndim == 2
# after
lhs_scale = lhs_scale[None, ...]  # rank 3

nd = lhs.ndim
ok = rhs.ndim == nd and (lhs_scale is None or lhs_scale.ndim == nd) and (rhs_scale is None or rhs_scale.ndim == nd)
assert ok

def same_rank(lhs, rhs, ls, rs) -> bool:
    n = lhs.ndim
    return rhs.ndim == n and (ls is None or ls.ndim == n) and (rs is None or rs.ndim == n)

When it happens

Trigger: lhs of shape (B,M,K) with rhs_scale of shape (K,N) (rank 2 vs 3); passing 2D matrices when a batched 3D call with scales is expected.

Common situations: Mixing batched and unbatched operands; scales saved as lower-rank arrays from a checkpoint; adding scales to previously-working dot code with mismatched rank.

Related errors

• process_id and num_processes must be nonnegative, with proce
• broadcast_in_dim target broadcast shape must have equal or h
• {side} contracting dim {i} of size {size} must be divisible
• 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

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