jax-ml/jax · error · TypeError

shardings should container 4 inputs, but got {len(shardings)

Error message

shardings should container 4 inputs, but got {len(shardings)}

What it means

Internal partitioning check in the cuDNN scaled matmul (block-scaled FP8) StableHLO path: the sharding-spec tuple passed around must contain exactly 4 entries (lhs, rhs, and two others such as scale/ouput specs). Getting a different count means the sharding metadata for the operator was constructed inconsistently.

Source

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

#   ([B], M, K1), ([B], N, K2)
# We define the following rule to apply necessary AllGather based on
# "Input specs", and to define the "Output spec".
# 1. If K1 == K2 != None and N == None:
#   - Input spec : ([B], M, K1), ([B], None, K2)
#   - 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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Align jax and jax-cudnn / nvidia plugin versions (pip install -U jax jax-cudnn)
2. If you build shardings yourself, ensure exactly 4 specs are passed (lhs, rhs, plus the two remaining specs)
3. Reproduce without custom sharding constraints to confirm the default path works, then re-add constraints incrementally

assert len(shardings) == 4, f'expected 4 shardings, got {len(shardings)}'

When it happens

Trigger: Running jax.lax.dot_general-based scaled matmul (mxfp8/nxfp8 paths) under a SPMD sharding where the collected shardings tuple for the operator does not have 4 elements — typically from custom partitioning logic or mismatched jax versions between jax and jax-cudnn.

Common situations: Version mismatch between jax and nvidia/jax-cudnn packages; writing custom sharding rules that build the shardings tuple by hand; passing 2D (unbatched) specs where 3D are expected elsewhere.

Related errors

• shardings specs rank should be 3, but got lhs: {len(lhs.spec
• shardings spec for batch dim should be same, but got lhs: {l
• callbacks are only supported in spmd computations when all m
• callbacks do not support specifying sharding inside spmd com
• {name} cannot accept args which are unreduced. Got {a.str_sh

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