{"record":{"id":"7c02423bd5dfb0ea","repo":"jax-ml/jax","slug":"dot-general-requires-lhs-batch-dimensions-and-rhs","errorCode":null,"errorMessage":"dot_general requires lhs batch dimensions and rhs batch dimensions to have the same shape, got {} and {}.","messagePattern":"dot_general requires lhs batch dimensions and rhs batch dimensions to have the same shape, got (.+?) and (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":5738,"sourceCode":"    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)\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(","sourceCodeStart":5720,"sourceCodeEnd":5756,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L5720-L5756","documentation":"Raised by lax.dot_general validation when the lhs and rhs batch dimensions do not have equal shapes: tuple(lhs.shape[i] for i in lhs_batch) != the corresponding rhs tuple. Paired batch axes must match exactly (or be definitely-equal under symbolic dims).","triggerScenarios":"jax.lax.dot_general with batch dims pairing axes of different sizes, e.g. lhs shape (2, 3, 4) batched with rhs shape (5, 4, 6) via (((), ()), ((0,), (0,))).","commonSituations":"Batched attention or block-matmul code where the batch axes of queries/keys drift apart (e.g. leading vs trailing batch axis, or extra vmap axes).","solutions":["Reshape/transpose so paired batch axes have identical sizes","Double-check you selected matching axes: often the batch axis of one operand is axis 0 and of the other axis 1","Use jax.vmap over the batch axis instead of explicit batch dimension numbers"],"exampleFix":"// before\nres = lax.dot_general(x, y, (((), ()), ((0,), (1,))))  # x.shape[0]=2 != y.shape[1]=5\n// after\ny = y.transpose(1, 0, 2)  # batch axis first\nres = lax.dot_general(x, y, (((), ()), ((0,), (0,))))","handlingStrategy":"validation","validationCode":"(lc, rc), (lb, rb) = dimension_numbers\nlhs_bshape = tuple(lhs.shape[i] for i in lb)\nrhs_bshape = tuple(rhs.shape[i] for i in rb)\nassert lhs_bshape == rhs_bshape, (lhs_bshape, rhs_bshape)","typeGuard":"def batch_shapes_match(lhs, rhs, dn):\n    _, (lb, rb) = dn\n    return tuple(lhs.shape[i] for i in lb) == tuple(rhs.shape[j] for j in rb)","tryCatchPattern":null,"preventionTips":["Standardize on leading batch axes across operands","Use jax.vmap(lax.dot_general) for batched cases"],"tags":["jax","dot-general","shape-mismatch","batch-dims"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}