jax-ml/jax · error · TypeError

dot_general requires lhs batch dimensions and rhs batch dime

Error message

dot_general requires lhs batch dimensions and rhs batch dimensions to have the same shape, got {} and {}.

What it means

Raised by lax.dot_general validation when the lhs and rhs batch dimensions do not have equal shapes: tuple(lhs.shape[i] for i in lhs_batch) != the corresponding rhs tuple. Paired batch axes must match exactly (or be definitely-equal under symbolic dims).

Source

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

    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)
  lhs_contract_or_batch = tuple(sorted(tuple(lhs_contracting) + tuple(lhs_batch)))
  lhs_tensored_shape = tuple_delete(lhs_shape, lhs_contract_or_batch)
  rhs_group = ()
  if isinstance(dimension_numbers, RaggedDotDimensionNumbers):
    rhs_group = tuple(dimension_numbers.rhs_group_dimensions)
  rhs_contract_or_batch_or_group = tuple(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Reshape/transpose so paired batch axes have identical sizes
  2. Double-check you selected matching axes: often the batch axis of one operand is axis 0 and of the other axis 1
  3. Use jax.vmap over the batch axis instead of explicit batch dimension numbers

Example fix

// before
res = lax.dot_general(x, y, (((), ()), ((0,), (1,))))  # x.shape[0]=2 != y.shape[1]=5
// after
y = y.transpose(1, 0, 2)  # batch axis first
res = lax.dot_general(x, y, (((), ()), ((0,), (0,))))
Defensive patterns

Strategy: validation

Validate before calling

(lc, rc), (lb, rb) = dimension_numbers
lhs_bshape = tuple(lhs.shape[i] for i in lb)
rhs_bshape = tuple(rhs.shape[i] for i in rb)
assert lhs_bshape == rhs_bshape, (lhs_bshape, rhs_bshape)

Type guard

def batch_shapes_match(lhs, rhs, dn):
    _, (lb, rb) = dn
    return tuple(lhs.shape[i] for i in lb) == tuple(rhs.shape[j] for j in rb)

Prevention

When it happens

Trigger: jax.lax.dot_general with batch dims pairing axes of different sizes, e.g. lhs shape (2, 3, 4) batched with rhs shape (5, 4, 6) via (((), ()), ((0,), (0,))).

Common situations: Batched attention or block-matmul code where the batch axes of queries/keys drift apart (e.g. leading vs trailing batch axis, or extra vmap axes).

Related errors

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