{"record":{"id":"9c5b5290ebbf9943","repo":"jax-ml/jax","slug":"shardings-should-container-4-inputs-but-got-len","errorCode":null,"errorMessage":"shardings should container 4 inputs, but got {len(shardings)}","messagePattern":"shardings should container 4 inputs, but got (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/cudnn/scaled_matmul_stablehlo.py","lineNumber":220,"sourceCode":"#   ([B], M, K1), ([B], N, K2)\n# We define the following rule to apply necessary AllGather based on\n# \"Input specs\", and to define the \"Output spec\".\n# 1. If K1 == K2 != None and N == None:\n#   - Input spec : ([B], M, K1), ([B], None, K2)\n#   - Output spec: ([B], M, None) -> AllReduce -> ([B], M, None)\n# 2. If K1 == K2 != None and M == N != None:\n#   - Input spec : ([B], M, K1), ([B], None, K2)\n#   - Output spec: ([B], M, None) -> ReduceScatter -> ([B], M, N)\n# 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","sourceCodeStart":202,"sourceCodeEnd":238,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/cudnn/scaled_matmul_stablehlo.py#L202-L238","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Align jax and jax-cudnn / nvidia plugin versions (pip install -U jax jax-cudnn)","If you build shardings yourself, ensure exactly 4 specs are passed (lhs, rhs, plus the two remaining specs)","Reproduce without custom sharding constraints to confirm the default path works, then re-add constraints incrementally"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"assert len(shardings) == 4, f'expected 4 shardings, got {len(shardings)}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pin jax and jax-cudnn versions together","Keep custom partitioning code covered by multi-device tests"],"tags":["jax","fp8","sharding","spmd","matmul"],"backgroundTag":"sharding-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}