{"record":{"id":"f8bdd9981dbd0caf","repo":"jax-ml/jax","slug":"dot-general-requires-contracting-dimensions-to-hav","errorCode":null,"errorMessage":"dot_general requires contracting dimensions to have the same shape, got {} and {}.","messagePattern":"dot_general requires contracting dimensions to have the same shape, got (.+?) and (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":5744,"sourceCode":"    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)\n  lhs_contract_or_batch = tuple(sorted(tuple(lhs_contracting) + tuple(lhs_batch)))\n  lhs_tensored_shape = tuple_delete(lhs_shape, lhs_contract_or_batch)\n  rhs_group = ()\n  if isinstance(dimension_numbers, RaggedDotDimensionNumbers):\n    rhs_group = tuple(dimension_numbers.rhs_group_dimensions)\n  rhs_contract_or_batch_or_group = tuple(\n      sorted(tuple(rhs_contracting) + tuple(rhs_batch) + rhs_group)\n  )\n  rhs_tensored_shape = tuple_delete(rhs_shape, rhs_contract_or_batch_or_group)\n  return batch_shape + lhs_tensored_shape + rhs_tensored_shape\n\ndef _dot_general_sharding_rule(lhs, rhs, *, dimension_numbers, precision,","sourceCodeStart":5726,"sourceCodeEnd":5762,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L5726-L5762","documentation":"Raised by lax.dot_general validation when paired contracting dimensions differ in size: lhs.shape[i] for each lhs contracting dim must equal the paired rhs.shape[j]. These are the axes summed over in the inner product, so they must match.","triggerScenarios":"jax.lax.dot_general where the K dimension of lhs != K of rhs, e.g. matmul-like call with (((), ())), ((1,), (0,)) on shapes (3, 4) @ (5, 3).","commonSituations":"Classic matmul inner-dimension mismatch surfacing through dot_general in jnp.matmul/jnp.einsum lowering; dynamic shapes where one operand was sliced differently.","solutions":["Fix operand shapes so contracting dims match (transpose or slice)","Check the pairing order of (lhs_contracting, rhs_contracting) tuples","Add an assert on lhs.shape[c] == rhs.shape[c] before the call in tests"],"exampleFix":"# before\nout = lax.dot_general(a, b, (((1,), (0,)), ((), ())))  # a: (3,4), b: (5,3)\n# after\nout = lax.dot_general(a, b.T, (((1,), (0,)), ((), ())))  # b.T: (3,5)","handlingStrategy":"validation","validationCode":"(lc, rc), _ = dimension_numbers\nfor i, j in zip(lc, rc):\n    assert lhs.shape[i] == rhs.shape[j], (i, j, lhs.shape[i], rhs.shape[j])","typeGuard":"def contracting_dims_match(lhs, rhs, dn):\n    (lc, rc), _ = dn\n    return all(lhs.shape[i] == rhs.shape[j] for i, j in zip(lc, rc))","tryCatchPattern":null,"preventionTips":["Assert inner dims equal in shape-checking helpers before jitted calls","Log shapes when this fires inside jit traces"],"tags":["jax","dot-general","shape-mismatch","contraction"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}