jax-ml/jax · error · TypeError

shardings spec for batch dim should be same, but got lhs: {l

Error message

shardings spec for batch dim should be same, but got lhs: {lhs.spec[0]} and rhs: {rhs.spec[0]}

What it means

The lhs and rhs sharding specs for the cuDNN scaled matmul must shard the batch dimension identically; this error reports a mismatch. (Note: the message string is missing its f-prefix, so it prints the literal braces — a known cosmetic bug that makes it look odd in logs.)

Source

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

# 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):
  _, _, 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 n_spec == None

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make spec[0] identical for lhs and rhs (usually both replicated or both sharded on the batch axis)
2. If operands genuinely have different batching, reshard or reshape before the matmul
3. After fixing, also confirm ranks are 3 to avoid the follow-on checks

Example fix

# before
lhs = with_sharding_constraint(lhs, NamedSharding(mesh, P('data', None, None)))
rhs = with_sharding_constraint(rhs, NamedSharding(mesh, P(None, None, 'model')))
# after
lhs = with_sharding_constraint(lhs, NamedSharding(mesh, P('data', None, None)))
rhs = with_sharding_constraint(rhs, NamedSharding(mesh, P('data', None, None)))

assert lhs_spec[0] == rhs_spec[0], 'batch-dim sharding must match between lhs and rhs'

When it happens

Trigger: Constraining lhs to a replicated batch dim while rhs is sharded on dim 0 (or vice versa) when calling block-scaled matmul under SPMD; e.g. with jax.lax.with_sharding_constraint where the two specs' first entries differ.

Common situations: Copy-pasting sharding constraints between operands of different layouts; partially updating specs when switching from local to multi-host input pipelines.

Related errors

• shardings should container 4 inputs, but got {len(shardings)
• shardings specs rank should be 3, but got lhs: {len(lhs.spec
• Only support preferred_element_type in (f32, bf16, f16), but
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m

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