{"record":{"id":"2c8072ee7b5be1a0","repo":"jax-ml/jax","slug":"ragged-dot-general-requires-the-group-count-last","errorCode":null,"errorMessage":"ragged_dot_general requires the group count (last dimension of group_sizes) to be static in Mode 1 (non-contracting) and Mode 2 (contracting).","messagePattern":"ragged_dot_general requires the group count \\(last dimension of group_sizes\\) to be static in Mode 1 \\(non-contracting\\) and Mode 2 \\(contracting\\)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":6408,"sourceCode":"    raise TypeError('expected rank of group_sizes to be >=1.')\n  if group_sizes.ndim != 1:\n    # Construct the expected shape [b...,x...,g] of group_sizes.\n    prefix_dims = _ragged_dot_prefix_dims(\n        mode, lhs.ndim, lhs_ragged_dim, lhs_batch, lhs_contracting\n    )\n    expected_gs_shape = tuple(lhs.shape[i] for i in prefix_dims)\n    expected_gs_shape += (group_sizes.shape[-1],)\n    # TODO(pravnar): Permit other broadcastable shapes.\n    if not core.definitely_equal_shape(group_sizes.shape, expected_gs_shape):\n      raise TypeError(\n          'expected group_sizes to have shape '\n          f'{expected_gs_shape}, got {group_sizes.shape}.'\n      )\n  num_groups = group_sizes.shape[-1]\n  if (mode in (RaggedDotMode.RAGGED_CONTRACTING,\n               RaggedDotMode.RAGGED_NONCONTRACTING)\n      and core.is_symbolic_dim(num_groups)):\n    raise TypeError(\n        'ragged_dot_general requires the group count (last dimension of '\n        'group_sizes) to be static in Mode 1 (non-contracting) and Mode 2 '\n        '(contracting).'\n    )\n\n  # Validate properties of the rhs group dimension(s).\n  rhs_group_dims = ragged_dot_dimension_numbers.rhs_group_dimensions\n  match mode:\n    case RaggedDotMode.RAGGED_CONTRACTING | RaggedDotMode.RAGGED_BATCH:\n      if len(rhs_group_dims) != 0:\n        raise TypeError(\n            'ragged_dot_general requires zero group dimensions in the rhs '\n            'when lhs ragged dimension is contracting or batch.'\n        )\n    case RaggedDotMode.RAGGED_NONCONTRACTING:\n      if len(rhs_group_dims) != 1:\n        raise TypeError(\n            'ragged_dot_general requires exactly one rhs group dimension '","sourceCodeStart":6390,"sourceCodeEnd":6426,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L6390-L6426","documentation":"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.","triggerScenarios":"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.","commonSituations":"Dynamic-shape pipelines exporting to XLA/IREE; converting MoE grouped-matmul code where num_groups came from a traced value.","solutions":["Make num_groups static: compute it in Python and reshape group_sizes with concrete g","Restructure to Mode 3 (ragged batch) if a dynamic group count is genuinely required","Avoid dynamic-shape wrappers around the group axis for this op"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"from jax._src import core\nnum_groups = group_sizes.shape[-1]\nassert not core.is_symbolic_dim(num_groups) or mode == 3, 'num_groups must be static in Mode 1/2'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Compute group counts in Python outside jit","Keep the group axis out of dynamic-shape regions; restructure to Mode 3 if group count must be dynamic"],"tags":["jax","ragged-dot-general","dynamic-shapes","jit"],"backgroundTag":"static-shape-required","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}