jax-ml/jax · error · TypeError

expected rhs group dimension size to be {num_groups}, got {r

Error message

expected rhs group dimension size to be {num_groups}, got {rhs.shape[rhs_group_dim]}.

What it means

The size of the rhs group dimension must equal the number of groups implied by group_sizes (len(group_sizes)). A mismatch means lhs and rhs disagree on how many groups the ragged dot partitions into.

Source

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

        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,
      rhs,
      dimension_numbers=ragged_dot_dimension_numbers,
      precision=precision,
      preferred_element_type=preferred_element_type,
      out_sharding=None,
  )
  if mode == RaggedDotMode.RAGGED_CONTRACTING:
    out_shape = (num_groups,) + out_shape
  return out_shape


def _ragged_dot_general_dtype_rule(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make group_sizes length equal rhs.shape[rhs_group_dim] (e.g. recompute routing counts for the actual number of experts)
2. Check for accidental slicing/padding of the expert weight stack
3. Verify num_experts config matches the loaded checkpoint's weight shape

Example fix

// before
group_sizes = jnp.array([2,3,1,4])        # 4 groups
w = jnp.zeros((8, 6, 16))                # 8 experts -> mismatch
// after
w = jnp.zeros((4, 6, 16))               # group dim == len(group_sizes)

assert rhs.shape[dn.rhs_group_dimensions[0]] == len(group_sizes)

def groups_match(rhs, dn, gs) -> bool:
    return rhs.shape[dn.rhs_group_dimensions[0]] == gs.shape[0]

When it happens

Trigger: Calling ragged_dot_general in RAGGED_NONCONTRACTING mode where rhs.shape[rhs_group_dim] != len(group_sizes). E.g. 8 experts in group_sizes but a weights tensor with leading dim 4.

Common situations: MoE routing: number of experts in the router/group_sizes doesn't match the stacked expert weight tensor's group dimension; off-by-one or stale weights after changing expert count config.

Related errors

• ragged_dot_general requires zero group dimensions in the rhs
• ragged_dot_general requires exactly one rhs group dimension
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac

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