jax-ml/jax · error · TypeError

ragged_dot_general requires zero group dimensions in the rhs

Error message

ragged_dot_general requires zero group dimensions in the rhs when lhs ragged dimension is contracting or batch.

What it means

ragged_dot_general validates the rhs group dimensions based on the RaggedDotMode. When the lhs ragged dimension is contracting or batch, the rhs must have zero group dimensions; the raggedness is captured entirely by lhs + group_sizes, so any rhs group dims are invalid.

Source

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

          'expected group_sizes to have shape '
          f'{expected_gs_shape}, got {group_sizes.shape}.'
      )
  num_groups = group_sizes.shape[-1]
  if (mode in (RaggedDotMode.RAGGED_CONTRACTING,
               RaggedDotMode.RAGGED_NONCONTRACTING)
      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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove rhs_group_dimensions from your RaggedDotDimensionNumbers (set it to ()) when using RAGGED_CONTRACTING or RAGGED_BATCH
2. If you intended grouped rhs (e.g. per-expert weights), use RaggedDotMode.RAGGED_NONCONTRACTING instead
3. Check that the ragged dimension placement on lhs matches the mode you chose

Example fix

// before
dn = RaggedDotDimensionNumbers((1,),(0,),(0,),(0,),(0,))  # rhs_group=(0,)
out = ragged_dot_general(x, w, gs, dn, mode=RaggedDotMode.RAGGED_CONTRACTING)
// after
dn = RaggedDotDimensionNumbers((1,),(0,),(0,),(0,))  # no rhs_group
out = ragged_dot_general(x, w, gs, dn, mode=RaggedDotMode.RAGGED_CONTRACTING)

dn = ...  # RaggedDotDimensionNumbers
if dn.rhs_group_dimensions and mode in (RaggedDotMode.RAGGED_CONTRACTING, RaggedDotMode.RAGGED_BATCH):
    dn = dn._replace(rhs_group_dimensions=())  # or raise

def valid_ragged_dn(dn, mode) -> bool:
    g = len(dn.rhs_group_dimensions)
    if mode is RaggedDotMode.RAGGED_NONCONTRACTING:
        return g == 1
    return g == 0

except TypeError as e: assert 'zero group dimensions' in str(e)

When it happens

Trigger: Calling jax.lax.ragged_dot_general (or jax.lax.experimental.ragged_dot) with RaggedDotMode.RAGGED_CONTRACTING or RAGGED_BATCH while passing a RaggedDotDimensionNumbers whose rhs_group_dimensions is non-empty.

Common situations: Copying a config from a noncontracting-mode example and switching the mode without clearing rhs_group_dimensions; misunderstanding that group dims on rhs are only allowed in noncontracting mode (e.g. MoE grouped-matmul with grouped weights).

Related errors

• ragged_dot_general requires exactly one rhs group dimension
• 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/1a36c7690d525c8c. Report an issue: GitHub.