jax-ml/jax · error · TypeError

shardings specs rank should be 3, but got lhs: {len(lhs.spec

Error message

shardings specs rank should be 3, but got lhs: {len(lhs.spec)} and rhs: {len(rhs.spec)}

What it means

The cuDNN scaled matmul partitioner requires the lhs and rhs sharding specs to each be rank 3 ([B, M/N, K]). This error fires when either spec has a rank other than 3, meaning the operands are not being viewed as batched 3D matrices as the kernel requires.

Source

Thrown at jax/_src/cudnn/scaled_matmul_stablehlo.py:225

#   - Output spec: ([B], M, None) -> AllReduce -> ([B], M, None)
# 2. If K1 == K2 != None and M == N != None:
#   - Input spec : ([B], M, K1), ([B], None, K2)
#   - Output spec: ([B], M, None) -> ReduceScatter -> ([B], M, N)
# 3. If N == M:
#   - Input specs : ([B], M, None), ([B], None, None)
#   - Output specs: ([B], M, None)
# 4. If N != M:
#   - Input spec : ([B], M, None), ([B], N, None)
#   - Output spec: ([B], M, N)
def _check_shardings(shardings):
  if len(shardings) != 4:
    msg = f"shardings should container 4 inputs, but got {len(shardings)}"
    raise TypeError(msg)
  lhs, rhs, _, _ = shardings
  if len(lhs.spec) != 3 or len(rhs.spec) != 3:
    msg = (f'shardings specs rank should be 3, but got lhs: {len(lhs.spec)} '
            'and rhs: {len(rhs.spec)}')
    raise TypeError(msg)
  if lhs.spec[0] != rhs.spec[0]:
    msg = ('shardings spec for batch dim should be same, but got lhs: '
            '{lhs.spec[0]} and rhs: {rhs.spec[0]}')
    raise TypeError(msg)


def _enable_reduce_scatter(lhs, rhs):
  _, m_spec, lhs_k_spec = lhs.spec
  _, n_spec, rhs_k_spec = rhs.spec
  return (
      lhs_k_spec != None
      and lhs_k_spec == rhs_k_spec
      and m_spec != None
      and m_spec == n_spec
  )


def _enable_all_reduce(lhs, rhs):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape operands to rank 3 with an explicit batch dimension (e.g. x[None, :, :])
2. Verify your custom sharding specs are length-3 tuples
3. Check jax/jax-cudnn version compatibility if you changed nothing custom

Example fix

// before
out = scaled_dot(lhs_2d, rhs_2d)
// after
out = scaled_dot(lhs_2d[None], rhs_2d[None])[0]

assert lhs.ndim == 3 and rhs.ndim == 3, 'scaled matmul operands must be rank 3 [B, M, K]/[B, K, N]'

When it happens

Trigger: Supplying 2D matrices (no batch dim) or higher-rank operands to the block-scaled matmul path; or sharding specs whose leading batch dims were contracted away, e.g. after reshaping operands before the scaled dot product.

Common situations: Feeding unbatched [M,K]@[K,N] inputs where [1,M,K]@[1,K,N] is required; custom GSPMD annotations with rank-2 specs; mismatches after operand reshapes.

Related errors

• shardings should container 4 inputs, but got {len(shardings)
• shardings spec for batch dim should be same, but got lhs: {l
• Only support preferred_element_type in (f32, bf16, f16), but
• {}: arrays must have same number of dimensions, got {}.
• Mapped away dimension of inputs passed to vmap should be sha

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