jax-ml/jax · error · TypeError

expected rank of group_sizes to be >=1.

Error message

expected rank of group_sizes to be >=1.

What it means

Shape validation for ragged_dot_general's group_sizes: a 0-D (scalar) group_sizes was passed, but rank >= 1 is required. group_sizes must at least be a 1-D vector of per-group sizes [g].

Source

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

          f'got {dim} for {arg_name} of rank {rank}.'
      )

  # Validate the lhs ragged dimension, and find out which mode we're in.
  if len(ragged_dot_dimension_numbers.lhs_ragged_dimensions) != 1:
    raise TypeError(
        'ragged_dot_general expects exactly one lhs ragged dimension.'
    )
  lhs_ragged_dim = ragged_dot_dimension_numbers.lhs_ragged_dimensions[0]
  _check_in_range(lhs_ragged_dim, lhs.ndim, 'lhs ragged dimension', 'lhs')
  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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the scalar into a 1-D array: group_sizes=jnp.array([n])
2. For batched/nested cases, give group_sizes shape [b..., x..., g] matching the ragged-dot prefix dims

Example fix

# before
out = lax.ragged_dot_general(x, y, jnp.array(10), dims)
# after
out = lax.ragged_dot_general(x, y, jnp.array([10]), dims)

assert group_sizes.ndim >= 1, 'group_sizes must be at least 1-D'
if np.ndim(group_sizes) == 0:
    group_sizes = jnp.array([group_sizes])

def valid_group_sizes(gs):
    return getattr(gs, 'ndim', np.ndim(gs)) >= 1

When it happens

Trigger: jax.lax.ragged_dot_general(..., group_sizes=jnp.array(5)) or a Python int/0-D tracer for group_sizes.

Common situations: Single-group use where the developer passes a scalar total instead of [total]; forgetting jnp.array wrapping around a list.

Related errors

• expected group_sizes to have shape {expected_gs_shape}, got
• 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/23866ea8f38cc3de. Report an issue: GitHub.