jax-ml/jax · error · TypeError

dot_general requires lhs batch dimensions to be disjoint fro

Error message

dot_general requires lhs batch dimensions to be disjoint from contracting dimensions, got lhs_batch {} and lhs_contracting {}.

What it means

Raised by lax.dot_general validation when the same index appears in both lhs_batch and lhs_contracting for the lhs operand. An axis cannot simultaneously be broadcast (batched) and summed (contracted).

Source

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

    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 "
           "shape, got {} and {}.")
    raise TypeError(msg.format(lhs_contracting_shape, rhs_contracting_shape))

  return _dot_general_shape_computation(lhs.shape, rhs.shape, dimension_numbers)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Decide whether the shared lhs axis should batch or contract, and remove it from the other list
2. Recompute dimension numbers from the intended einsum with jnp.einsum(..., out='opt_einsum') or by inspection of the summation labels

Example fix

// before
res = lax.dot_general(x, y, (((1,), (0,)), ((1,), (0,))))  # axis 1 both batched and contracted
// after: axis 1 is contracting only
res = lax.dot_general(x, y, (((1,), (0,)), ((), ())))

(lhs_c, _), (lhs_b, _) = dimension_numbers
assert not (set(lhs_c) & set(lhs_b)), 'axis both batched and contracted on lhs'

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 with dimension_numbers where set(dimension_numbers[0][0]) & set(dimension_numbers[1][0]) is non-empty, e.g. (((1,), ()), ((1,), ())).

Common situations: Converting a tensordot/einsum expression to dot_general and reusing an axis index in both roles; ambiguous contraction specs copied from other frameworks.

Related errors

• dot_general requires rhs 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/23ff53bf3222ef1b. Report an issue: GitHub.