jax-ml/jax · error · TypeError

dot_general requires contracting dimensions to have the same

Error message

dot_general requires contracting dimensions to have the same shape, got {} and {}.

What it means

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.

Source

Thrown at jax/_src/lax/lax.py:5744

    msg = ("dot_general requires lhs batch dimensions to be disjoint from "
           "contracting dimensions, got lhs_batch {} and lhs_contracting {}.")
    raise TypeError(msg.format(lhs_batch, lhs_contracting))
  if rhs_contracting_set & rhs_batch_set:
    msg = ("dot_general requires rhs batch dimensions to be disjoint from "
           "contracting dimensions, got rhs_batch {} and rhs_contracting {}.")
    raise TypeError(msg.format(rhs_batch, rhs_contracting))
  lhs_batch_shape = tuple(lhs.shape[i] for i in lhs_batch)
  rhs_batch_shape = tuple(rhs.shape[i] for i in rhs_batch)
  if not core.definitely_equal_shape(lhs_batch_shape, rhs_batch_shape):
    msg = ("dot_general requires lhs batch dimensions and rhs batch dimensions "
           "to have the same shape, got {} and {}.")
    raise TypeError(msg.format(lhs_batch_shape, rhs_batch_shape))
  lhs_contracting_shape = tuple(lhs.shape[i] for i in lhs_contracting)
  rhs_contracting_shape = tuple(rhs.shape[i] for i in rhs_contracting)
  if not core.definitely_equal_shape(lhs_contracting_shape, rhs_contracting_shape):
    msg = ("dot_general requires contracting dimensions to have the same "
           "shape, got {} and {}.")
    raise TypeError(msg.format(lhs_contracting_shape, rhs_contracting_shape))

  return _dot_general_shape_computation(lhs.shape, rhs.shape, dimension_numbers)

def _dot_general_shape_computation(lhs_shape, rhs_shape, dimension_numbers):
  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = _from_maybe_ragged(dimension_numbers)
  batch_shape = tuple(lhs_shape[i] for i in lhs_batch)
  lhs_contract_or_batch = tuple(sorted(tuple(lhs_contracting) + tuple(lhs_batch)))
  lhs_tensored_shape = tuple_delete(lhs_shape, lhs_contract_or_batch)
  rhs_group = ()
  if isinstance(dimension_numbers, RaggedDotDimensionNumbers):
    rhs_group = tuple(dimension_numbers.rhs_group_dimensions)
  rhs_contract_or_batch_or_group = tuple(
      sorted(tuple(rhs_contracting) + tuple(rhs_batch) + rhs_group)
  )
  rhs_tensored_shape = tuple_delete(rhs_shape, rhs_contract_or_batch_or_group)
  return batch_shape + lhs_tensored_shape + rhs_tensored_shape

def _dot_general_sharding_rule(lhs, rhs, *, dimension_numbers, precision,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Fix operand shapes so contracting dims match (transpose or slice)
2. Check the pairing order of (lhs_contracting, rhs_contracting) tuples
3. Add an assert on lhs.shape[c] == rhs.shape[c] before the call in tests

Example fix

# before
out = lax.dot_general(a, b, (((1,), (0,)), ((), ())))  # a: (3,4), b: (5,3)
# after
out = lax.dot_general(a, b.T, (((1,), (0,)), ((), ())))  # b.T: (3,5)

(lc, rc), _ = dimension_numbers
for i, j in zip(lc, rc):
    assert lhs.shape[i] == rhs.shape[j], (i, j, lhs.shape[i], rhs.shape[j])

def contracting_dims_match(lhs, rhs, dn):
    (lc, rc), _ = dn
    return all(lhs.shape[i] == rhs.shape[j] for i, j in zip(lc, rc))

When it happens

Trigger: 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).

Common situations: Classic matmul inner-dimension mismatch surfacing through dot_general in jnp.matmul/jnp.einsum lowering; dynamic shapes where one operand was sliced differently.

Related errors

• dot_general requires lhs contracting dimensions to be distin
• dot_general requires rhs contracting dimensions to be distin
• dot_general requires lhs batch dimensions to be disjoint fro
• dot_general requires rhs batch dimensions to be disjoint fro
• dot_general requires lhs batch dimensions and rhs batch dime

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/f8bdd9981dbd0caf. Report an issue: GitHub.