{"record":{"id":"0bf7a16945b630e0","repo":"jax-ml/jax","slug":"shardings-spec-for-batch-dim-should-be-same-but-g","errorCode":null,"errorMessage":"shardings spec for batch dim should be same, but got lhs: {lhs.spec[0]} and rhs: {rhs.spec[0]}","messagePattern":"shardings spec for batch dim should be same, but got lhs: (.+?) and rhs: (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/cudnn/scaled_matmul_stablehlo.py","lineNumber":229,"sourceCode":"# 3. If N == M:\n#   - Input specs : ([B], M, None), ([B], None, None)\n#   - Output specs: ([B], M, None)\n# 4. If N != M:\n#   - Input spec : ([B], M, None), ([B], N, None)\n#   - Output spec: ([B], M, N)\ndef _check_shardings(shardings):\n  if len(shardings) != 4:\n    msg = f\"shardings should container 4 inputs, but got {len(shardings)}\"\n    raise TypeError(msg)\n  lhs, rhs, _, _ = shardings\n  if len(lhs.spec) != 3 or len(rhs.spec) != 3:\n    msg = (f'shardings specs rank should be 3, but got lhs: {len(lhs.spec)} '\n            'and rhs: {len(rhs.spec)}')\n    raise TypeError(msg)\n  if lhs.spec[0] != rhs.spec[0]:\n    msg = ('shardings spec for batch dim should be same, but got lhs: '\n            '{lhs.spec[0]} and rhs: {rhs.spec[0]}')\n    raise TypeError(msg)\n\n\ndef _enable_reduce_scatter(lhs, rhs):\n  _, m_spec, lhs_k_spec = lhs.spec\n  _, n_spec, rhs_k_spec = rhs.spec\n  return (\n      lhs_k_spec != None\n      and lhs_k_spec == rhs_k_spec\n      and m_spec != None\n      and m_spec == n_spec\n  )\n\n\ndef _enable_all_reduce(lhs, rhs):\n  _, _, lhs_k_spec = lhs.spec\n  _, n_spec, rhs_k_spec = rhs.spec\n  return lhs_k_spec != None and lhs_k_spec == rhs_k_spec and n_spec == None\n","sourceCodeStart":211,"sourceCodeEnd":247,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/cudnn/scaled_matmul_stablehlo.py#L211-L247","documentation":"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.)","triggerScenarios":"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.","commonSituations":"Copy-pasting sharding constraints between operands of different layouts; partially updating specs when switching from local to multi-host input pipelines.","solutions":["Make spec[0] identical for lhs and rhs (usually both replicated or both sharded on the batch axis)","If operands genuinely have different batching, reshard or reshape before the matmul","After fixing, also confirm ranks are 3 to avoid the follow-on checks"],"exampleFix":"# before\nlhs = with_sharding_constraint(lhs, NamedSharding(mesh, P('data', None, None)))\nrhs = with_sharding_constraint(rhs, NamedSharding(mesh, P(None, None, 'model')))\n# after\nlhs = with_sharding_constraint(lhs, NamedSharding(mesh, P('data', None, None)))\nrhs = with_sharding_constraint(rhs, NamedSharding(mesh, P('data', None, None)))","handlingStrategy":"validation","validationCode":"assert lhs_spec[0] == rhs_spec[0], 'batch-dim sharding must match between lhs and rhs'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive both specs from one shared layout object","Review sharding constraints in code review when meshes change"],"tags":["jax","fp8","sharding","batch-dim","matmul"],"backgroundTag":"sharding-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}