jax-ml/jax · error · TypeError

ragged_dot_general requires the group count (last dimension

Error message

ragged_dot_general requires the group count (last dimension of group_sizes) to be static in Mode 1 (non-contracting) and Mode 2 (contracting).

What it means

In ragged_dot_general Mode 1 (ragged non-contracting) and Mode 2 (ragged contracting), the number of groups — the last dimension of group_sizes — must be a static integer, not a symbolic/dynamic dimension. Only Mode 3 (ragged batch) may have a dynamic group count.

Source

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

    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:
        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 '

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make num_groups static: compute it in Python and reshape group_sizes with concrete g
2. Restructure to Mode 3 (ragged batch) if a dynamic group count is genuinely required
3. Avoid dynamic-shape wrappers around the group axis for this op

from jax._src import core
num_groups = group_sizes.shape[-1]
assert not core.is_symbolic_dim(num_groups) or mode == 3, 'num_groups must be static in Mode 1/2'

When it happens

Trigger: Calling ragged_dot_general in Mode 1/2 under jax.jit with dynamic shapes (e.g. from jax.experimental.dynamic_shape or layout/shape polymorphism) so group_sizes.shape[-1] is a symbolic dim.

Common situations: Dynamic-shape pipelines exporting to XLA/IREE; converting MoE grouped-matmul code where num_groups came from a traced value.

Related errors

• hessenberg requires the last dimension of a to be constant,
• Formatting arguments to checkify.check need to be PyTrees of
• Value of type {type(self)} is not convertible to float.
• Value of type {type(self)} is not convertible to complex.
• {self.__class__.__name__} has no attribute {name}

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