{"record":{"id":"6637836904d79b4e","repo":"jax-ml/jax","slug":"shardings-specs-rank-should-be-3-but-got-lhs-le","errorCode":null,"errorMessage":"shardings specs rank should be 3, but got lhs: {len(lhs.spec)} and rhs: {len(rhs.spec)}","messagePattern":"shardings specs rank should be 3, but got lhs: (.+?) and rhs: (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/cudnn/scaled_matmul_stablehlo.py","lineNumber":225,"sourceCode":"#   - 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\n      and m_spec == n_spec\n  )\n\n\ndef _enable_all_reduce(lhs, rhs):","sourceCodeStart":207,"sourceCodeEnd":243,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/cudnn/scaled_matmul_stablehlo.py#L207-L243","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Reshape operands to rank 3 with an explicit batch dimension (e.g. x[None, :, :])","Verify your custom sharding specs are length-3 tuples","Check jax/jax-cudnn version compatibility if you changed nothing custom"],"exampleFix":"// before\nout = scaled_dot(lhs_2d, rhs_2d)\n// after\nout = scaled_dot(lhs_2d[None], rhs_2d[None])[0]","handlingStrategy":"validation","validationCode":"assert lhs.ndim == 3 and rhs.ndim == 3, 'scaled matmul operands must be rank 3 [B, M, K]/[B, K, N]'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Standardize on 3D operands at API boundaries","Assert ranks in data-prep code"],"tags":["jax","fp8","sharding","rank-mismatch","matmul"],"backgroundTag":"tensor-rank-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}