jax-ml/jax · error · TypeError

dot_general requires lhs contracting dimensions to be distin

Error message

dot_general requires lhs contracting dimensions to be distinct, got lhs_contracting {lhs_contracting}.

What it means

Raised by lax.dot_general validation when lhs_contracting contains duplicate indices. Contracting dims are summed over; listing one twice is ambiguous and illegal.

Source

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

    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 "
           "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))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove duplicates from lhs_contracting; to contract two lhs axes you need two distinct rhs axes
2. Verify len(lhs_contracting) == len(rhs_contracting)
3. Let jnp.einsum lower to dot_general automatically for multi-axis contractions

Example fix

// before
res = lax.dot_general(x, y, (((1, 1), (0, 0)), ((), ())))
// after: contract distinct axes
res = lax.dot_general(x, y, (((1, 2), (0, 0)), ((), ())))

lhs_c = dimension_numbers[0][0]
assert len(set(lhs_c)) == len(lhs_c), 'duplicate lhs contracting dims'

def valid_contracting(dn):
    (lc, rc), _ = dn
    return len(set(lc)) == len(lc) and len(set(rc)) == len(rc) and len(lc) == len(rc)

When it happens

Trigger: jax.lax.dot_general with dimension_numbers[0][0] repeating an index, e.g. (((1, 1), (0, 0)), ((), ())) intending a double contraction.

Common situations: Trying to contract two axes of the same operand against one axis of the other; porting einsum specs like 'iji,j->' naively into dimension numbers.

Related errors

• 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 dimension numbers to be nonnegative
• dot_general requires rhs dimension numbers to be nonnegative

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