jax-ml/jax · error · TypeError

dot_general requires lhs dimension numbers to be nonnegative

Error message

dot_general requires lhs dimension numbers to be nonnegative and less than the number of axes of the lhs value, got lhs_batch of {lhs_batch} and lhs_contracting of {lhs_contracting} for lhs of rank {lhs.ndim}

What it means

dot_general takes dimension_numbers = ((lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch)). All lhs dimension indices must be in [0, lhs.ndim). Negative or too-large indices (numpy-style negative indexing is not supported) raise TypeError with the offending lists and rank.

Source

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

  preferred_bitwidth = np.dtype(preferred_element_type).itemsize
  if preferred_bitwidth < input_bitwidth:
    raise TypeError("`preferred_element_type` must not be narrower than the "
                    "original type.")


def _dot_general_shape_rule(lhs, rhs, *, dimension_numbers, precision,
                            preferred_element_type: DTypeLike | None,
                            out_sharding):
  if out_sharding is not None and not isinstance(out_sharding, NamedSharding):
    raise NotImplementedError
  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = _from_maybe_ragged(dimension_numbers)
  if not all(np.all(np.greater_equal(d, 0)) and np.all(np.less(d, lhs.ndim))
             for d in (lhs_contracting, lhs_batch)):
    msg = ("dot_general requires lhs dimension numbers to be nonnegative and "
           "less than the number of axes of the lhs value, got "
           f"lhs_batch of {lhs_batch} and lhs_contracting of {lhs_contracting} "
           f"for lhs of rank {lhs.ndim}")
    raise TypeError(msg)
  if not all(np.all(np.greater_equal(d, 0)) and np.all(np.less(d, rhs.ndim))
             for d in (rhs_contracting, rhs_batch)):
    msg = ("dot_general requires rhs dimension numbers to be nonnegative and "
           "less than the number of axes of the rhs value, got "
           f"rhs_batch of {rhs_batch} and rhs_contracting of {rhs_contracting} "
           f"for rhs of rank {rhs.ndim}")
    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):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use nonnegative axis indices within range for lhs contracting and batch dims
2. Remember structure ((lhs_contract, rhs_contract), (lhs_batch, rhs_batch)) and verify each list
3. Prefer jnp.einsum/jnp.tensordot/jnp.matmul which accept friendlier axis specs
4. Validate indices against lhs.ndim before calling

Example fix

// before
out = lax.dot_general(a, b, ((-1,), (0,)), ((), ()))

// after
out = lax.dot_general(a, b, ((a.ndim - 1,), (0,)), ((), ()))

assert all(0 <= d < lhs.ndim for d in (*lhs_contracting, *lhs_batch)), 'bad lhs dims'

def valid_lhs_dims(dn, lhs) -> bool:
    (lc, _), (lb, _) = dn
    return all(0 <= d < lhs.ndim for d in (*lc, *lb))

When it happens

Trigger: lax.dot_general(a, b, ((-1,), ()), (...)) using a negative axis; passing contracting dim 3 for a rank-2 lhs; misordered tuple where batch dims land in the contracting slot.

Common situations: Porting numpy einsum or jnp.tensordot axis lists that allow negative indices; misreading the nested dimension_numbers tuple structure; off-by-one axis constants after refactoring.

Related errors

• dot_general requires rhs dimension numbers to be nonnegative
• 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/fcf4223fc79666db. Report an issue: GitHub.