{"record":{"id":"ed690a9a6ab20073","repo":"jax-ml/jax","slug":"dot-general-requires-lhs-batch-dimensions-to-be-di","errorCode":null,"errorMessage":"dot_general requires lhs batch dimensions to be distinct, got lhs_batch {lhs_batch}.","messagePattern":"dot_general requires lhs batch dimensions to be distinct, got lhs_batch (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":5712,"sourceCode":"           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)\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 \"","sourceCodeStart":5694,"sourceCodeEnd":5730,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L5694-L5730","documentation":"Raised by lax.dot_general's dimension-number validation when the lhs_batch sequence contains duplicate dimension indices. Batch dimensions must each be distinct because each one maps a separate lhs axis to a separate rhs axis for broadcasting.","triggerScenarios":"Calling jax.lax.dot_general(lhs, rhs, dimension_numbers=...) where dimension_numbers[1][0] (lhs batch dims) repeats an index, e.g. (((), ()), ((0, 0), (0, 0))).","commonSituations":"Hand-constructing dimension_numbers tuples for batched matmuls instead of using jnp.einsum or lax.batch_matmul; typos when copying dimension numbers from XLA/HLO dumps.","solutions":["Remove duplicate entries from lhs_batch so each lhs batch dimension appears once","Prefer higher-level APIs (jnp.matmul, jnp.einsum, lax.batch_matmul) that build dimension_numbers for you","Print/validate dimension_numbers with sets before calling dot_general"],"exampleFix":"// before\nres = lax.dot_general(x, y, (((), ()), ((0, 0), (1, 1))))\n// after\nres = lax.dot_general(x, y, (((), ()), ((0, 1), (0, 1))))","handlingStrategy":"validation","validationCode":"lhs_batch = dimension_numbers[1][0]\nassert len(set(lhs_batch)) == len(lhs_batch), 'duplicate lhs batch dims'","typeGuard":"def valid_batch_dims(dn):\n    (lc, rc), (lb, rb) = dn\n    return len(set(lb)) == len(lb) and len(set(rb)) == len(rb) and len(lb) == len(rb)","tryCatchPattern":null,"preventionTips":["Build dimension numbers next to the shapes they reference","Prefer jnp.einsum/jnp.tensordot which construct valid dimension numbers"],"tags":["jax","dot-general","dimension-numbers","batch-dims"],"backgroundTag":"invalid-dimension-numbers","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}