{"record":{"id":"6d29ce548df95405","repo":"jax-ml/jax","slug":"ragged-dot-general-requires-exactly-one-rhs-group","errorCode":null,"errorMessage":"ragged_dot_general requires exactly one rhs group dimension when lhs ragged dimension is noncontracting.","messagePattern":"ragged_dot_general requires exactly one rhs group dimension when lhs ragged dimension is noncontracting\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":6425,"sourceCode":"      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 '\n            'when lhs ragged dimension is noncontracting.'\n        )\n      rhs_group_dim = rhs_group_dims[0]\n      _check_in_range(rhs_group_dim, rhs.ndim, 'rhs group dimension', 'rhs')\n      if rhs_group_dim in rhs_batch or rhs_group_dim in rhs_contracting:\n        raise TypeError(\n            'ragged_dot_general requires rhs group dimension numbers to be '\n            'distinct from contracting and batch dimensions.'\n        )\n      if rhs.shape[rhs_group_dim] != num_groups:\n        raise TypeError(\n            'expected rhs group dimension size to be '\n            f'{num_groups}, got {rhs.shape[rhs_group_dim]}.'\n        )\n\n  out_shape = _dot_general_shape_rule(\n      lhs,","sourceCodeStart":6407,"sourceCodeEnd":6443,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L6407-L6443","documentation":"In RAGGED_NONCONTRACTING mode, ragged_dot_general requires exactly one rhs group dimension so the rhs can be partitioned into the same number of groups as the lhs ragged dimension. Zero or multiple group dims on rhs fail this check.","triggerScenarios":"Calling ragged_dot_general with mode=RaggedDotMode.RAGGED_NONCONTRACTING where rhs_group_dimensions has length != 1 (typically 0 because it was left empty).","commonSituations":"Grouped MoE-style matmul where the weight matrix has a leading num_groups dimension, but the user omitted rhs_group_dimensions when constructing the dimension numbers.","solutions":["Add exactly one rhs group dimension index (the dim of rhs holding the groups, e.g. 0 for stacked expert weights) to RaggedDotDimensionNumbers","Verify that dim size equals len(group_sizes) (num_groups)","If your rhs has no group dim, you probably want RAGGED_CONTRACTING/RAGGED_BATCH mode instead"],"exampleFix":"// before\ndn = RaggedDotDimensionNumbers((2,),(1,),(0,),(0,))  # missing rhs group\nout = ragged_dot_general(x, w, gs, dn, mode=RaggedDotMode.RAGGED_NONCONTRACTING)\n// after\ndn = RaggedDotDimensionNumbers((2,),(1,),(0,),(0,),(0,))  # rhs_group=(0,)\nout = ragged_dot_general(x, w, gs, dn, mode=RaggedDotMode.RAGGED_NONCONTRACTING)","handlingStrategy":"validation","validationCode":"assert len(dn.rhs_group_dimensions) == 1 or mode is not RaggedDotMode.RAGGED_NONCONTRACTING","typeGuard":"def has_one_rhs_group(dn, mode) -> bool:\n    return mode is not RaggedDotMode.RAGGED_NONCONTRACTING or len(dn.rhs_group_dimensions) == 1","tryCatchPattern":null,"preventionTips":["Keep dimension-number construction next to mode selection in one function"],"tags":["jax","ragged-dot","shape-validation","moe"],"backgroundTag":"invalid-dimension-configuration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}