jax-ml/jax · error · TypeError

expected group_sizes to have shape {expected_gs_shape}, got

Error message

expected group_sizes to have shape {expected_gs_shape}, got {group_sizes.shape}.

What it means

When group_sizes has rank > 1, ragged_dot_general requires its shape to be exactly the broadcast prefix (batch and non-contracting dims of lhs) plus the trailing group axis, i.e. [b..., x..., g]; anything else fails this check.

Source

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

  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
  )

  # Validate the shape of group_sizes, if it is something other than [g].
  if group_sizes.ndim == 0:
    raise TypeError('expected rank of group_sizes to be >=1.')
  if group_sizes.ndim != 1:
    # Construct the expected shape [b...,x...,g] of group_sizes.
    prefix_dims = _ragged_dot_prefix_dims(
        mode, lhs.ndim, lhs_ragged_dim, lhs_batch, lhs_contracting
    )
    expected_gs_shape = tuple(lhs.shape[i] for i in prefix_dims)
    expected_gs_shape += (group_sizes.shape[-1],)
    # TODO(pravnar): Permit other broadcastable shapes.
    if not core.definitely_equal_shape(group_sizes.shape, expected_gs_shape):
      raise TypeError(
          '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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Transpose/reshape group_sizes so the group axis is last and lhs batch/non-contracting dims come first in lhs order
2. Broadcast a base (b..., 1) group-size array with jnp.broadcast_to to the expected (b..., g) shape

Example fix

# before
out = lax.ragged_dot_general(x, y, gs, dims)  # gs.shape == (g, b)
# after
gs = gs.T  # (b, g)
out = lax.ragged_dot_general(x, y, gs, dims)

if group_sizes.ndim > 1:
    prefix = tuple(lhs.shape[i] for i in prefix_dims)
    expected = prefix + (group_sizes.shape[-1],)
    assert group_sizes.shape == expected, (group_sizes.shape, expected)
# or broadcast defensively:
# group_sizes = jnp.broadcast_to(group_sizes, expected)

def group_sizes_shape_ok(lhs, gs, prefix_dims):
    if gs.ndim <= 1: return True
    return gs.shape == tuple(lhs.shape[i] for i in prefix_dims) + (gs.shape[-1],)

When it happens

Trigger: Passing group_sizes shaped (g, b) instead of (b, g), or omitting lhs batch dims from group_sizes, in a batched ragged_dot_general call.

Common situations: Nested/moe-style grouped GEMMs where per-batch group sizes were stacked in the wrong axis order; forgetting to repeat group sizes along a non-contracting lhs axis.

Related errors

• expected rank of group_sizes to be >=1.
• 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
• Incorrect output shape for return value #{i}: Expected: {out

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