jax-ml/jax · error · TypeError

dot_general requires equal numbers of lhs_batch and rhs_batc

Error message

dot_general requires equal numbers of lhs_batch and rhs_batch dimensions, got lhs_batch {} and rhs_batch {}.

What it means

dot_general batch semantics require len(lhs_batch) == len(rhs_batch): each batch dim on the lhs pairs with one on the rhs. Unequal batch lists make the output shape undefined, so TypeError is raised with both lists.

Source

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

  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = _from_maybe_ragged(dimension_numbers)
  if not all(np.all(np.greater_equal(d, 0)) and np.all(np.less(d, lhs.ndim))
             for d in (lhs_contracting, lhs_batch)):
    msg = ("dot_general requires lhs dimension numbers to be nonnegative and "
           "less than the number of axes of the lhs value, got "
           f"lhs_batch of {lhs_batch} and lhs_contracting of {lhs_contracting} "
           f"for lhs of rank {lhs.ndim}")
    raise TypeError(msg)
  if not all(np.all(np.greater_equal(d, 0)) and np.all(np.less(d, rhs.ndim))
             for d in (rhs_contracting, rhs_batch)):
    msg = ("dot_general requires rhs dimension numbers to be nonnegative and "
           "less than the number of axes of the rhs value, got "
           f"rhs_batch of {rhs_batch} and rhs_contracting of {rhs_contracting} "
           f"for rhs of rank {rhs.ndim}")
    raise TypeError(msg)
  if len(lhs_batch) != len(rhs_batch):
    msg = ("dot_general requires equal numbers of lhs_batch and rhs_batch "
           "dimensions, got lhs_batch {} and rhs_batch {}.")
    raise TypeError(msg.format(lhs_batch, rhs_batch))
  lhs_contracting_set, lhs_batch_set = set(lhs_contracting), set(lhs_batch)
  rhs_contracting_set, rhs_batch_set = set(rhs_contracting), set(rhs_batch)
  if len(lhs_batch_set) != len(lhs_batch):
    msg = ("dot_general requires lhs batch dimensions to be distinct, got "
           f"lhs_batch {lhs_batch}.")
    raise TypeError(msg)
  if len(rhs_batch_set) != len(rhs_batch):
    msg = ("dot_general requires rhs batch dimensions to be distinct, got "
           f"rhs_batch {rhs_batch}.")
    raise TypeError(msg)
  if len(lhs_contracting_set) != len(lhs_contracting):
    msg = ("dot_general requires lhs contracting dimensions to be distinct, "
           f"got lhs_contracting {lhs_contracting}.")
    raise TypeError(msg)
  if len(rhs_contracting_set) != len(rhs_contracting):
    msg = ("dot_general requires rhs contracting dimensions to be distinct, "
           f"got rhs_contracting {rhs_contracting}.")
    raise TypeError(msg)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Supply matching-length lhs_batch and rhs_batch lists
2. For an einsum like 'ii->i', both sides batch on the same index: (((), ()), ((0,), (0,)))
3. Double-check tuple structure: dimension_numbers = ((lhs_contract, rhs_contract), (lhs_batch, rhs_batch))
4. Prefer jnp.einsum for diag-style patterns

Example fix

// before
out = lax.dot_general(a, b, ((0,), (1,)), ((0,), ()))  # rhs_batch empty

// after
out = lax.dot_general(a, b, ((0,), (1,)), ((0,), (0,)))

assert len(lhs_batch) == len(rhs_batch), 'batch dim count mismatch'

def batch_dims_paired(dn) -> bool:
    (_, _), (lb, rb) = dn
    return len(lb) == len(rb)

When it happens

Trigger: lax.dot_general(a, b, ((0,), ())), ((), (0,))) style calls where one side lists batch dims and the other doesn't; e.g. ((0,), (1,)), ((0,), ()) — lhs_batch of length 1, rhs_batch of length 0.

Common situations: Misnesting the dimension_numbers tuple so a batch list ends up empty on one side; porting einsum 'ii->' patterns that batch on one operand only; hand-writing batched contractions.

Related errors

• dot_general requires lhs batch dimensions to be distinct, go
• dot_general requires rhs batch dimensions to be distinct, go
• dot_general requires lhs batch dimensions to be disjoint fro
• dot_general requires rhs batch dimensions to be disjoint fro
• dot_general requires lhs dimension numbers to be nonnegative

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