{"record":{"id":"0aa77e8a0e01f530","repo":"jax-ml/jax","slug":"dot-general-requires-equal-numbers-of-lhs-batch-an","errorCode":null,"errorMessage":"dot_general requires equal numbers of lhs_batch and rhs_batch dimensions, got lhs_batch {} and rhs_batch {}.","messagePattern":"dot_general requires equal numbers of lhs_batch and rhs_batch dimensions, got lhs_batch (.+?) and rhs_batch (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":5706,"sourceCode":"  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = _from_maybe_ragged(dimension_numbers)\n  if not all(np.all(np.greater_equal(d, 0)) and np.all(np.less(d, lhs.ndim))\n             for d in (lhs_contracting, lhs_batch)):\n    msg = (\"dot_general requires lhs dimension numbers to be nonnegative and \"\n           \"less than the number of axes of the lhs value, got \"\n           f\"lhs_batch of {lhs_batch} and lhs_contracting of {lhs_contracting} \"\n           f\"for lhs of rank {lhs.ndim}\")\n    raise TypeError(msg)\n  if not all(np.all(np.greater_equal(d, 0)) and np.all(np.less(d, rhs.ndim))\n             for d in (rhs_contracting, rhs_batch)):\n    msg = (\"dot_general requires rhs dimension numbers to be nonnegative and \"\n           \"less than the number of axes of the rhs value, got \"\n           f\"rhs_batch of {rhs_batch} and rhs_contracting of {rhs_contracting} \"\n           f\"for rhs of rank {rhs.ndim}\")\n    raise TypeError(msg)\n  if len(lhs_batch) != len(rhs_batch):\n    msg = (\"dot_general requires equal numbers of lhs_batch and rhs_batch \"\n           \"dimensions, got lhs_batch {} and rhs_batch {}.\")\n    raise TypeError(msg.format(lhs_batch, rhs_batch))\n  lhs_contracting_set, lhs_batch_set = set(lhs_contracting), set(lhs_batch)\n  rhs_contracting_set, rhs_batch_set = set(rhs_contracting), set(rhs_batch)\n  if len(lhs_batch_set) != len(lhs_batch):\n    msg = (\"dot_general requires lhs batch dimensions to be distinct, got \"\n           f\"lhs_batch {lhs_batch}.\")\n    raise TypeError(msg)\n  if len(rhs_batch_set) != len(rhs_batch):\n    msg = (\"dot_general requires rhs batch dimensions to be distinct, got \"\n           f\"rhs_batch {rhs_batch}.\")\n    raise TypeError(msg)\n  if len(lhs_contracting_set) != len(lhs_contracting):\n    msg = (\"dot_general requires lhs contracting dimensions to be distinct, \"\n           f\"got lhs_contracting {lhs_contracting}.\")\n    raise TypeError(msg)\n  if len(rhs_contracting_set) != len(rhs_contracting):\n    msg = (\"dot_general requires rhs contracting dimensions to be distinct, \"\n           f\"got rhs_contracting {rhs_contracting}.\")\n    raise TypeError(msg)","sourceCodeStart":5688,"sourceCodeEnd":5724,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L5688-L5724","documentation":"dot_general batch semantics require len(lhs_batch) == len(rhs_batch): each batch dim on the lhs pairs with one on the rhs. Unequal batch lists make the output shape undefined, so TypeError is raised with both lists.","triggerScenarios":"lax.dot_general(a, b, ((0,), ())), ((), (0,))) style calls where one side lists batch dims and the other doesn't; e.g. ((0,), (1,)), ((0,), ()) — lhs_batch of length 1, rhs_batch of length 0.","commonSituations":"Misnesting the dimension_numbers tuple so a batch list ends up empty on one side; porting einsum 'ii->' patterns that batch on one operand only; hand-writing batched contractions.","solutions":["Supply matching-length lhs_batch and rhs_batch lists","For an einsum like 'ii->i', both sides batch on the same index: (((), ()), ((0,), (0,)))","Double-check tuple structure: dimension_numbers = ((lhs_contract, rhs_contract), (lhs_batch, rhs_batch))","Prefer jnp.einsum for diag-style patterns"],"exampleFix":"// before\nout = lax.dot_general(a, b, ((0,), (1,)), ((0,), ()))  # rhs_batch empty\n\n// after\nout = lax.dot_general(a, b, ((0,), (1,)), ((0,), (0,)))","handlingStrategy":"validation","validationCode":"assert len(lhs_batch) == len(rhs_batch), 'batch dim count mismatch'","typeGuard":"def batch_dims_paired(dn) -> bool:\n    (_, _), (lb, rb) = dn\n    return len(lb) == len(rb)","tryCatchPattern":null,"preventionTips":["Always supply lhs_batch and rhs_batch as equal-length lists","For einsum 'ii' patterns batch on both operands at the same index","Build dimension_numbers via a tested helper instead of inline literals"],"tags":["jax","dot-general","batch-dims","dimension-numbers"],"backgroundTag":"batch-dim-count-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}