jax-ml/jax · error · TypeError

dot_general requires rhs batch dimensions to be disjoint fro

Error message

dot_general requires rhs batch dimensions to be disjoint from contracting dimensions, got rhs_batch {} and rhs_contracting {}.

What it means

Raised by lax.dot_general validation when the same index appears in both rhs_batch and rhs_contracting for the rhs operand. A rhs axis cannot be both a broadcast (batch) dimension and a contracted dimension.

Source

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the overlapping index from either rhs_batch or rhs_contracting depending on intent
2. Validate disjointness of the two sets before calling dot_general

Example fix

// before
res = lax.dot_general(q, k, (((2,), (1,)), ((), (2,))))
// after
res = lax.dot_general(q, k, (((2,), (1,)), ((), ())))

(_, rhs_c), (_, rhs_b) = dimension_numbers
assert not (set(rhs_c) & set(rhs_b)), 'axis both batched and contracted on rhs'

def disjoint_dims(dn):
    (lc, rc), (lb, rb) = dn
    return not (set(lc) & set(lb)) and not (set(rc) & set(rb))

When it happens

Trigger: jax.lax.dot_general where set(dimension_numbers[0][1]) & set(dimension_numbers[1][1]) is non-empty, e.g. (((0,), (2,)), ((), (2,))).

Common situations: Porting attention-style batched matmuls into explicit dimension numbers and accidentally tagging the feature/head axis as both batch and contract.

Related errors

• dot_general requires lhs batch dimensions to be disjoint fro
• dot_general requires equal numbers of lhs_batch and rhs_batc
• dot_general requires lhs batch dimensions to be distinct, go
• dot_general requires rhs batch dimensions to be distinct, go
• dot_general requires lhs contracting dimensions to be distin

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