jax-ml/jax · error · TypeError

ragged_dot_general requires {dim_name} numbers to be nonnega

Error message

ragged_dot_general requires {dim_name} numbers to be nonnegative and less than the number of axes of the {arg_name} value, got {dim} for {arg_name} of rank {rank}.

What it means

Range check inside ragged_dot_general: a supplied dimension index (lhs/rhs ragged dim, or rhs group dim) is negative or >= the rank of the corresponding operand.

Source

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

      return batch + contract[: contract.index(ragged_dim)]
    case RaggedDotMode.RAGGED_BATCH:
      return batch[: batch.index(ragged_dim)]


def _ragged_dot_general_shape_rule(
    lhs,
    rhs,
    group_sizes,
    *,
    ragged_dot_dimension_numbers,
    precision,
    preferred_element_type: DTypeLike | None,
    group_offset,
    out_sharding,
):
  def _check_in_range(dim, rank, dim_name, arg_name):
    if dim < 0 or dim >= rank:
      raise TypeError(
          f'ragged_dot_general requires {dim_name} numbers to be nonnegative '
          f'and less than the number of axes of the {arg_name} value, '
          f'got {dim} for {arg_name} of rank {rank}.'
      )

  # Validate the lhs ragged dimension, and find out which mode we're in.
  if len(ragged_dot_dimension_numbers.lhs_ragged_dimensions) != 1:
    raise TypeError(
        'ragged_dot_general expects exactly one lhs ragged dimension.'
    )
  lhs_ragged_dim = ragged_dot_dimension_numbers.lhs_ragged_dimensions[0]
  _check_in_range(lhs_ragged_dim, lhs.ndim, 'lhs ragged dimension', 'lhs')
  mode = _ragged_dot_mode(lhs.ndim, ragged_dot_dimension_numbers)

  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = (
      ragged_dot_dimension_numbers.dot_dimension_numbers
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use nonnegative indices strictly less than the operand's .ndim
2. Recompute all dim indices after any reshape/transpose of lhs/rhs
3. Add asserts on 0 <= dim < operand.ndim before the call

Example fix

# before
out = lax.ragged_dot_general(x, y, gs, dims_with_ragged_dim=-1)
# after
out = lax.ragged_dot_general(x, y, gs, dims_with_ragged_dim=x.ndim - 1)

def check_dim(dim, rank):
    assert 0 <= dim < rank, f'{dim} not in [0, {rank})'
check_dim(lhs_ragged_dim, lhs.ndim)
for d in rhs_group_dims: check_dim(d, rhs.ndim)

def dim_in_range(dim, operand):
    return isinstance(dim, int) and 0 <= dim < operand.ndim

When it happens

Trigger: jax.lax.ragged_dot_general where lhs_ragged_dim >= lhs.ndim or rhs_group_dim >= rhs.ndim (e.g. using -1 as an index is not allowed, and off-by-one on rank).

Common situations: Python-numpy habit of negative indices (not accepted here); shapes changed after adding a batch axis so old dim indices now exceed the rank.

Related errors

• dot_general requires lhs dimension numbers to be nonnegative
• dot_general requires rhs dimension numbers to be nonnegative
• lhs_ragged_dim {lhs_ragged_dim} not found in lhs_noncontract
• ragged_dot_general expects exactly one lhs ragged dimension.
• expected rank of group_sizes to be >=1.

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