{"record":{"id":"1a36c7690d525c8c","repo":"jax-ml/jax","slug":"ragged-dot-general-requires-zero-group-dimensions","errorCode":null,"errorMessage":"ragged_dot_general requires zero group dimensions in the rhs when lhs ragged dimension is contracting or batch.","messagePattern":"ragged_dot_general requires zero group dimensions in the rhs when lhs ragged dimension is contracting or batch\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":6419,"sourceCode":"          '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 '\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(","sourceCodeStart":6401,"sourceCodeEnd":6437,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L6401-L6437","documentation":"ragged_dot_general validates the rhs group dimensions based on the RaggedDotMode. When the lhs ragged dimension is contracting or batch, the rhs must have zero group dimensions; the raggedness is captured entirely by lhs + group_sizes, so any rhs group dims are invalid.","triggerScenarios":"Calling jax.lax.ragged_dot_general (or jax.lax.experimental.ragged_dot) with RaggedDotMode.RAGGED_CONTRACTING or RAGGED_BATCH while passing a RaggedDotDimensionNumbers whose rhs_group_dimensions is non-empty.","commonSituations":"Copying a config from a noncontracting-mode example and switching the mode without clearing rhs_group_dimensions; misunderstanding that group dims on rhs are only allowed in noncontracting mode (e.g. MoE grouped-matmul with grouped weights).","solutions":["Remove rhs_group_dimensions from your RaggedDotDimensionNumbers (set it to ()) when using RAGGED_CONTRACTING or RAGGED_BATCH","If you intended grouped rhs (e.g. per-expert weights), use RaggedDotMode.RAGGED_NONCONTRACTING instead","Check that the ragged dimension placement on lhs matches the mode you chose"],"exampleFix":"// before\ndn = RaggedDotDimensionNumbers((1,),(0,),(0,),(0,),(0,))  # rhs_group=(0,)\nout = ragged_dot_general(x, w, gs, dn, mode=RaggedDotMode.RAGGED_CONTRACTING)\n// after\ndn = RaggedDotDimensionNumbers((1,),(0,),(0,),(0,))  # no rhs_group\nout = ragged_dot_general(x, w, gs, dn, mode=RaggedDotMode.RAGGED_CONTRACTING)","handlingStrategy":"validation","validationCode":"dn = ...  # RaggedDotDimensionNumbers\nif dn.rhs_group_dimensions and mode in (RaggedDotMode.RAGGED_CONTRACTING, RaggedDotMode.RAGGED_BATCH):\n    dn = dn._replace(rhs_group_dimensions=())  # or raise","typeGuard":"def valid_ragged_dn(dn, mode) -> bool:\n    g = len(dn.rhs_group_dimensions)\n    if mode is RaggedDotMode.RAGGED_NONCONTRACTING:\n        return g == 1\n    return g == 0","tryCatchPattern":"except TypeError as e: assert 'zero group dimensions' in str(e)","preventionTips":["Build RaggedDotDimensionNumbers via a helper keyed on mode","Assert group-dim count matches mode in unit tests"],"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"}