jax-ml/jax · error · TypeError

dot_general requires lhs batch dimensions to be distinct, go

Error message

dot_general requires lhs batch dimensions to be distinct, got lhs_batch {lhs_batch}.

What it means

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.

Source

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

           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)
  if lhs_contracting_set & lhs_batch_set:
    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 "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove duplicate entries from lhs_batch so each lhs batch dimension appears once
2. Prefer higher-level APIs (jnp.matmul, jnp.einsum, lax.batch_matmul) that build dimension_numbers for you
3. Print/validate dimension_numbers with sets before calling dot_general

Example fix

// before
res = lax.dot_general(x, y, (((), ()), ((0, 0), (1, 1))))
// after
res = lax.dot_general(x, y, (((), ()), ((0, 1), (0, 1))))

lhs_batch = dimension_numbers[1][0]
assert len(set(lhs_batch)) == len(lhs_batch), 'duplicate lhs batch dims'

def valid_batch_dims(dn):
    (lc, rc), (lb, rb) = dn
    return len(set(lb)) == len(lb) and len(set(rb)) == len(rb) and len(lb) == len(rb)

When it happens

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

Common situations: 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.

Related errors

• dot_general requires equal numbers of lhs_batch and rhs_batc
• 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/ed690a9a6ab20073. Report an issue: GitHub.