jax-ml/jax · error · TypeError

ragged_dot_general requires exactly one rhs group dimension

Error message

ragged_dot_general requires exactly one rhs group dimension when lhs ragged dimension is noncontracting.

What it means

In RAGGED_NONCONTRACTING mode, ragged_dot_general requires exactly one rhs group dimension so the rhs can be partitioned into the same number of groups as the lhs ragged dimension. Zero or multiple group dims on rhs fail this check.

Source

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

      and core.is_symbolic_dim(num_groups)):
    raise TypeError(
        'ragged_dot_general requires the group count (last dimension of '
        'group_sizes) to be static in Mode 1 (non-contracting) and Mode 2 '
        '(contracting).'
    )

  # Validate properties of the rhs group dimension(s).
  rhs_group_dims = ragged_dot_dimension_numbers.rhs_group_dimensions
  match mode:
    case RaggedDotMode.RAGGED_CONTRACTING | RaggedDotMode.RAGGED_BATCH:
      if len(rhs_group_dims) != 0:
        raise TypeError(
            'ragged_dot_general requires zero group dimensions in the rhs '
            'when lhs ragged dimension is contracting or batch.'
        )
    case RaggedDotMode.RAGGED_NONCONTRACTING:
      if len(rhs_group_dims) != 1:
        raise TypeError(
            'ragged_dot_general requires exactly one rhs group dimension '
            'when lhs ragged dimension is noncontracting.'
        )
      rhs_group_dim = rhs_group_dims[0]
      _check_in_range(rhs_group_dim, rhs.ndim, 'rhs group dimension', 'rhs')
      if rhs_group_dim in rhs_batch or rhs_group_dim in rhs_contracting:
        raise TypeError(
            'ragged_dot_general requires rhs group dimension numbers to be '
            'distinct from contracting and batch dimensions.'
        )
      if rhs.shape[rhs_group_dim] != num_groups:
        raise TypeError(
            'expected rhs group dimension size to be '
            f'{num_groups}, got {rhs.shape[rhs_group_dim]}.'
        )

  out_shape = _dot_general_shape_rule(
      lhs,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Add exactly one rhs group dimension index (the dim of rhs holding the groups, e.g. 0 for stacked expert weights) to RaggedDotDimensionNumbers
2. Verify that dim size equals len(group_sizes) (num_groups)
3. If your rhs has no group dim, you probably want RAGGED_CONTRACTING/RAGGED_BATCH mode instead

Example fix

// before
dn = RaggedDotDimensionNumbers((2,),(1,),(0,),(0,))  # missing rhs group
out = ragged_dot_general(x, w, gs, dn, mode=RaggedDotMode.RAGGED_NONCONTRACTING)
// after
dn = RaggedDotDimensionNumbers((2,),(1,),(0,),(0,),(0,))  # rhs_group=(0,)
out = ragged_dot_general(x, w, gs, dn, mode=RaggedDotMode.RAGGED_NONCONTRACTING)

assert len(dn.rhs_group_dimensions) == 1 or mode is not RaggedDotMode.RAGGED_NONCONTRACTING

def has_one_rhs_group(dn, mode) -> bool:
    return mode is not RaggedDotMode.RAGGED_NONCONTRACTING or len(dn.rhs_group_dimensions) == 1

When it happens

Trigger: Calling ragged_dot_general with mode=RaggedDotMode.RAGGED_NONCONTRACTING where rhs_group_dimensions has length != 1 (typically 0 because it was left empty).

Common situations: Grouped MoE-style matmul where the weight matrix has a leading num_groups dimension, but the user omitted rhs_group_dimensions when constructing the dimension numbers.

Related errors

• ragged_dot_general requires zero group dimensions in the rhs
• ragged_dot_general requires rhs group dimension numbers to b
• expected rhs group dimension size to be {num_groups}, got {r
• scan got `length` argument of {} which disagrees with leadin
• conv_general_dilated batch_group_count must divide lhs batch

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