{"record":{"id":"cb71e1b998f374f9","repo":"jax-ml/jax","slug":"lhs-and-rhs-must-have-the-same-number-of-batch-dim","errorCode":null,"errorMessage":"LHS and RHS must have the same number of batch dimensions, got {len(lhs_batch)} and {len(rhs_batch)}.","messagePattern":"LHS and RHS must have the same number of batch dimensions, got (.+?) and (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/scaled_dot.py","lineNumber":78,"sourceCode":"  \"\"\"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):\n    batch_dims_sizes = [\n        lhs.shape[i_lhs],\n        rhs.shape[i_rhs],\n    ]\n    if lhs_scale is not None:\n      batch_dims_sizes.append(lhs_scale.shape[i_lhs])\n    if rhs_scale is not None:\n      batch_dims_sizes.append(rhs_scale.shape[i_rhs])","sourceCodeStart":60,"sourceCodeEnd":96,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/scaled_dot.py#L60-L96","documentation":"scaled_dot's dimension_numbers must declare the same number of batch dimensions for lhs and rhs. When `len(lhs_batch) != len(rhs_batch)` the validation raises TypeError with both counts.","triggerScenarios":"Passing dimension numbers where lhs has 1 batch dim and rhs has 0, e.g. ((0,), (), (2,), (1,)) with mismatched batch lists; forgetting the rhs batch dim in a custom dimension_numbers tuple.","commonSituations":"Hand-constructing lax-style dimension numbers ((lhs_contract, rhs_contract), (lhs_batch, rhs_batch)); adapting dot_general dimension numbers to the scaled variant and dropping one entry.","solutions":["Make the batch dimension lists equal length, typically both empty for plain matmuls or both containing the batch axis","Prefer passing dimension numbers as the standard ((lc, rc), (lb, rb)) tuple format and double-check symmetry","Test with the same shapes you use in an equivalent jnp.einsum to derive correct dims"],"exampleFix":"# before\ndims = ((2,), (1,), (0,), ())  # lhs has batch, rhs doesn't\n# after\ndims = ((2,), (1,), (0,), (0,))  # matching batch dims","handlingStrategy":"validation","validationCode":"assert len(lhs_batch) == len(rhs_batch), 'batch dim counts must match'\nassert len(lhs_contract) == len(rhs_contract)","typeGuard":"def valid_dim_numbers(dims) -> bool:\n    (lc, rc), (lb, rb) = dims\n    return len(lb) == len(rb) and len(lc) == len(rc)","tryCatchPattern":null,"preventionTips":["Derive dimension numbers from jnp.einsum or dot_general conventions; keep batch lists symmetric"],"tags":["jax","scaled-dot","dimension-numbers"],"backgroundTag":"dimension-numbers-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}