jax-ml/jax · error · TypeError

ragged_dot_general expects exactly one lhs ragged dimension.

Error message

ragged_dot_general expects exactly one lhs ragged dimension.

What it means

ragged_dot_general only supports exactly one ragged dimension on the lhs; this fires when len(lhs_ragged_dimensions) != 1 (zero or multiple entries).

Source

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

    group_sizes,
    *,
    ragged_dot_dimension_numbers,
    precision,
    preferred_element_type: DTypeLike | None,
    group_offset,
    out_sharding,
):
  def _check_in_range(dim, rank, dim_name, arg_name):
    if dim < 0 or dim >= rank:
      raise TypeError(
          f'ragged_dot_general requires {dim_name} numbers to be nonnegative '
          f'and less than the number of axes of the {arg_name} value, '
          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
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Supply exactly one lhs ragged dimension; model multiple ragged axes with separate ragged_dot_general calls or reshapes

assert len(rdn.lhs_ragged_dimensions) == 1, 'exactly one lhs ragged dim required'

def single_ragged_dim(rdn):
    return len(rdn.lhs_ragged_dimensions) == 1

When it happens

Trigger: Passing RaggedDotDimensionNumbers with lhs_ragged_dimensions=[] or with two indices, e.g. (1, 2), to jax.lax.ragged_dot_general.

Common situations: Migrating from XLA RaggedDot which may allow multiple ragged dims; misunderstanding that only the lhs carries one ragged dim (rhs raggedness comes via group dims).

Related errors

• lhs_ragged_dim {lhs_ragged_dim} not found in lhs_noncontract
• No 4+ dimensional dimension_number defaults.
• convolution dimension_numbers list/tuple must be length 3, g
• convolution dimension_numbers elements must be strings, got
• convolution dimension_numbers[{}] must have len equal to the

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