jax-ml/jax · error · TypeError

dot_general requires rhs contracting dimensions to be distin

Error message

dot_general requires rhs contracting dimensions to be distinct, got rhs_contracting {rhs_contracting}.

What it means

Raised by lax.dot_general validation when rhs_contracting contains duplicate indices. Each rhs contracting dimension must be distinct since each pairs with exactly one lhs contracting dimension for the sum.

Source

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Deduplicate rhs_contracting so each rhs axis appears once
2. Ensure the pairing (lhs_contracting[i], rhs_contracting[i]) is what you intend
3. Use jnp.einsum/jnp.tensordot for complex contractions

Example fix

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

rhs_c = dimension_numbers[0][1]
assert len(set(rhs_c)) == len(rhs_c), 'duplicate rhs contracting dims'

def valid_contracting(dn):
    (lc, rc), _ = dn
    return all(len(set(x)) == len(x) for x in (lc, rc))

When it happens

Trigger: jax.lax.dot_general with dimension_numbers[0][1] containing repeats, e.g. (((0, 0), (1, 1)), ((), ())).

Common situations: Hand-written transposition rules, custom_jvp/vjp rules, or ported XLA configs where the rhs contracting tuple was copy-pasted and edited incorrectly.

Related errors

• dot_general requires lhs 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/47ee0460a6c9141c. Report an issue: GitHub.