{"record":{"id":"2322dc02a3d2b118","repo":"jax-ml/jax","slug":"dot-general-requires-rhs-batch-dimensions-to-be-di-2322dc","errorCode":null,"errorMessage":"dot_general requires rhs batch dimensions to be disjoint from contracting dimensions, got rhs_batch {} and rhs_contracting {}.","messagePattern":"dot_general requires rhs batch dimensions to be disjoint from contracting dimensions, got rhs_batch (.+?) and rhs_contracting (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":5732,"sourceCode":"    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)\n  if lhs_contracting_set & lhs_batch_set:\n    msg = (\"dot_general requires lhs batch dimensions to be disjoint from \"\n           \"contracting dimensions, got lhs_batch {} and lhs_contracting {}.\")\n    raise TypeError(msg.format(lhs_batch, lhs_contracting))\n  if rhs_contracting_set & rhs_batch_set:\n    msg = (\"dot_general requires rhs batch dimensions to be disjoint from \"\n           \"contracting dimensions, got rhs_batch {} and rhs_contracting {}.\")\n    raise TypeError(msg.format(rhs_batch, rhs_contracting))\n  lhs_batch_shape = tuple(lhs.shape[i] for i in lhs_batch)\n  rhs_batch_shape = tuple(rhs.shape[i] for i in rhs_batch)\n  if not core.definitely_equal_shape(lhs_batch_shape, rhs_batch_shape):\n    msg = (\"dot_general requires lhs batch dimensions and rhs batch dimensions \"\n           \"to have the same shape, got {} and {}.\")\n    raise TypeError(msg.format(lhs_batch_shape, rhs_batch_shape))\n  lhs_contracting_shape = tuple(lhs.shape[i] for i in lhs_contracting)\n  rhs_contracting_shape = tuple(rhs.shape[i] for i in rhs_contracting)\n  if not core.definitely_equal_shape(lhs_contracting_shape, rhs_contracting_shape):\n    msg = (\"dot_general requires contracting dimensions to have the same \"\n           \"shape, got {} and {}.\")\n    raise TypeError(msg.format(lhs_contracting_shape, rhs_contracting_shape))\n\n  return _dot_general_shape_computation(lhs.shape, rhs.shape, dimension_numbers)\n\ndef _dot_general_shape_computation(lhs_shape, rhs_shape, dimension_numbers):\n  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = _from_maybe_ragged(dimension_numbers)\n  batch_shape = tuple(lhs_shape[i] for i in lhs_batch)","sourceCodeStart":5714,"sourceCodeEnd":5750,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L5714-L5750","documentation":"Raised by lax.dot_general validation when the same index appears in both rhs_batch and rhs_contracting for the rhs operand. A rhs axis cannot be both a broadcast (batch) dimension and a contracted dimension.","triggerScenarios":"jax.lax.dot_general where set(dimension_numbers[0][1]) & set(dimension_numbers[1][1]) is non-empty, e.g. (((0,), (2,)), ((), (2,))).","commonSituations":"Porting attention-style batched matmuls into explicit dimension numbers and accidentally tagging the feature/head axis as both batch and contract.","solutions":["Remove the overlapping index from either rhs_batch or rhs_contracting depending on intent","Validate disjointness of the two sets before calling dot_general"],"exampleFix":"// before\nres = lax.dot_general(q, k, (((2,), (1,)), ((), (2,))))\n// after\nres = lax.dot_general(q, k, (((2,), (1,)), ((), ())))","handlingStrategy":"validation","validationCode":"(_, rhs_c), (_, rhs_b) = dimension_numbers\nassert not (set(rhs_c) & set(rhs_b)), 'axis both batched and contracted on rhs'","typeGuard":"def disjoint_dims(dn):\n    (lc, rc), (lb, rb) = dn\n    return not (set(lc) & set(lb)) and not (set(rc) & set(rb))","tryCatchPattern":null,"preventionTips":["Write per-operand axis role maps in comments next to dimension_numbers"],"tags":["jax","dot-general","dimension-numbers","batch-dims","contraction"],"backgroundTag":"invalid-dimension-numbers","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}