{"record":{"id":"8c3a9405ded2583b","repo":"jax-ml/jax","slug":"all-input-tensors-must-have-the-same-rank-got-lhs","errorCode":null,"errorMessage":"All input tensors must have the same rank. Got lhs rank: {lhs.ndim} rhs rank: {rhs.ndim} lhs_scale rank: {lhs_scale.ndim if lhs_scale is not None else 'N/A'} rhs_scale rank: {rhs_scale.ndim if rhs_scale is not None else 'N/A'}.","messagePattern":"All input tensors must have the same rank\\. Got lhs rank: (.+?) rhs rank: (.+?) lhs_scale rank: (.+?) rhs_scale rank: (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/scaled_dot.py","lineNumber":70,"sourceCode":"    lhs: Array,\n    rhs: Array,\n    lhs_scale: Array | None,\n    rhs_scale: Array | None,\n    *,\n    dimension_numbers: lax.DotDimensionNumbers,\n    preferred_element_type: DTypeLike | None,\n):\n  \"\"\"Validates the inputs to scaled_dot.\"\"\"\n  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = dimension_numbers\n\n  ndims = [lhs.ndim, rhs.ndim]\n  if lhs_scale is not None:\n    ndims.append(lhs_scale.ndim)\n  if rhs_scale is not None:\n    ndims.append(rhs_scale.ndim)\n\n  if max(ndims) != min(ndims):\n    raise TypeError(\n        \"All input tensors must have the same rank. Got lhs rank:\"\n        f\" {lhs.ndim} rhs rank: {rhs.ndim} lhs_scale rank:\"\n        f\" {lhs_scale.ndim if lhs_scale is not None else 'N/A'} rhs_scale\"\n        f\" rank: {rhs_scale.ndim if rhs_scale is not None else 'N/A'}.\"\n    )\n\n  if len(lhs_batch) != len(rhs_batch):\n    raise TypeError(\n        \"LHS and RHS must have the same number of batch dimensions, got\"\n        f\" {len(lhs_batch)} and {len(rhs_batch)}.\"\n    )\n  if len(lhs_contracting) != len(rhs_contracting):\n    raise TypeError(\n        \"LHS and RHS must have the same number of contracting dimensions, got\"\n        f\" {len(lhs_contracting)} and {len(rhs_contracting)}.\"\n    )\n\n  for i_lhs, i_rhs in zip(lhs_batch, rhs_batch):","sourceCodeStart":52,"sourceCodeEnd":88,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/scaled_dot.py#L52-L88","documentation":"scaled_dot requires lhs, rhs and any provided lhs_scale/rhs_scale to all have the same rank (number of dimensions). `_scaled_dot_validate_inputs` compares min/max of the collected ndims and raises TypeError listing each rank when they differ.","triggerScenarios":"lhs of shape (B,M,K) with rhs_scale of shape (K,N) (rank 2 vs 3); passing 2D matrices when a batched 3D call with scales is expected.","commonSituations":"Mixing batched and unbatched operands; scales saved as lower-rank arrays from a checkpoint; adding scales to previously-working dot code with mismatched rank.","solutions":["Expand dims so every input has the same rank (e.g. scale[None] for batched operands)","Reshape/unsqueeze operands or scales consistently","Verify all tensors' .ndim right before the call in tests"],"exampleFix":"# before\nlhs.ndim == 3, lhs_scale.ndim == 2\n# after\nlhs_scale = lhs_scale[None, ...]  # rank 3","handlingStrategy":"type-guard","validationCode":"nd = lhs.ndim\nok = rhs.ndim == nd and (lhs_scale is None or lhs_scale.ndim == nd) and (rhs_scale is None or rhs_scale.ndim == nd)\nassert ok","typeGuard":"def same_rank(lhs, rhs, ls, rs) -> bool:\n    n = lhs.ndim\n    return rhs.ndim == n and (ls is None or ls.ndim == n) and (rs is None or rs.ndim == n)","tryCatchPattern":null,"preventionTips":["Expand dims on scales when batching operands","Assert .ndim equality in unit tests"],"tags":["jax","scaled-dot","rank-validation"],"backgroundTag":"rank-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}