{"record":{"id":"5c7fdbfb3863dc0a","repo":"jax-ml/jax","slug":"dot-general-requires-rhs-dimension-numbers-to-be-n","errorCode":null,"errorMessage":"dot_general requires rhs dimension numbers to be nonnegative and less than the number of axes of the rhs value, got rhs_batch of {rhs_batch} and rhs_contracting of {rhs_contracting} for rhs of rank {rhs.ndim}","messagePattern":"dot_general requires rhs dimension numbers to be nonnegative and less than the number of axes of the rhs value, got rhs_batch of (.+?) and rhs_contracting of (.+?) for rhs of rank (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":5702,"sourceCode":"                            preferred_element_type: DTypeLike | None,\n                            out_sharding):\n  if out_sharding is not None and not isinstance(out_sharding, NamedSharding):\n    raise NotImplementedError\n  (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)","sourceCodeStart":5684,"sourceCodeEnd":5720,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L5684-L5720","documentation":"Same bounds check as 917 but for the rhs operand: every dimension index in rhs_contracting and rhs_batch must satisfy 0 <= d < rhs.ndim. Negative indices (not supported) or indices at/above the rank raise TypeError echoing the lists and rhs rank.","triggerScenarios":"lax.dot_general with rhs contracting dim equal to rhs.ndim, or a negative rhs batch index; swapping lhs/rhs dimension lists so rhs gets lhs-range indices.","commonSituations":"Asymmetric tensor contractions where lhs and rhs have different ranks; copy-paste of dimension_numbers between calls with different operands; negative-axis habits from numpy.","solutions":["Clamp rhs indices to [0, rhs.ndim)","Ensure the rhs lists are in the second position of each pair: ((lhs_c, rhs_c), (lhs_b, rhs_b))","Use tensordot/einsum for readability when hand-building contractions","Print rhs.ndim and the dimension_numbers to spot the offending index"],"exampleFix":"// before\nout = lax.dot_general(a, b, ((1,), (2,)), ((), ()))  # b has rank 2\n\n// after\nout = lax.dot_general(a, b, ((1,), (1,)), ((), ()))","handlingStrategy":"validation","validationCode":"assert all(0 <= d < rhs.ndim for d in (*rhs_contracting, *rhs_batch)), 'bad rhs dims'","typeGuard":"def valid_rhs_dims(dn, rhs) -> bool:\n    (_, rc), (_, rb) = dn\n    return all(0 <= d < rhs.ndim for d in (*rc, *rb))","tryCatchPattern":null,"preventionTips":["Remember dimension_numbers structure ((lhs_c, rhs_c), (lhs_b, rhs_b))","Validate rhs indices separately when ranks differ between operands","Log rank + dimension_numbers on contraction bugs"],"tags":["jax","dot-general","dimension-numbers","index-out-of-range"],"backgroundTag":"axis-index-out-of-range","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}