{"record":{"id":"ca939f80259eb453","repo":"jax-ml/jax","slug":"expected-rhs-group-dimension-size-to-be-num-group","errorCode":null,"errorMessage":"expected rhs group dimension size to be {num_groups}, got {rhs.shape[rhs_group_dim]}.","messagePattern":"expected rhs group dimension size to be (.+?), got (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":6437,"sourceCode":"        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,\n      rhs,\n      dimension_numbers=ragged_dot_dimension_numbers,\n      precision=precision,\n      preferred_element_type=preferred_element_type,\n      out_sharding=None,\n  )\n  if mode == RaggedDotMode.RAGGED_CONTRACTING:\n    out_shape = (num_groups,) + out_shape\n  return out_shape\n\n\ndef _ragged_dot_general_dtype_rule(","sourceCodeStart":6419,"sourceCodeEnd":6455,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L6419-L6455","documentation":"The size of the rhs group dimension must equal the number of groups implied by group_sizes (len(group_sizes)). A mismatch means lhs and rhs disagree on how many groups the ragged dot partitions into.","triggerScenarios":"Calling ragged_dot_general in RAGGED_NONCONTRACTING mode where rhs.shape[rhs_group_dim] != len(group_sizes). E.g. 8 experts in group_sizes but a weights tensor with leading dim 4.","commonSituations":"MoE routing: number of experts in the router/group_sizes doesn't match the stacked expert weight tensor's group dimension; off-by-one or stale weights after changing expert count config.","solutions":["Make group_sizes length equal rhs.shape[rhs_group_dim] (e.g. recompute routing counts for the actual number of experts)","Check for accidental slicing/padding of the expert weight stack","Verify num_experts config matches the loaded checkpoint's weight shape"],"exampleFix":"// before\ngroup_sizes = jnp.array([2,3,1,4])        # 4 groups\nw = jnp.zeros((8, 6, 16))                # 8 experts -> mismatch\n// after\nw = jnp.zeros((4, 6, 16))               # group dim == len(group_sizes)","handlingStrategy":"validation","validationCode":"assert rhs.shape[dn.rhs_group_dimensions[0]] == len(group_sizes)","typeGuard":"def groups_match(rhs, dn, gs) -> bool:\n    return rhs.shape[dn.rhs_group_dimensions[0]] == gs.shape[0]","tryCatchPattern":null,"preventionTips":["Single-source num_experts from config into both router and weights","Validate expert-count consistency at model init"],"tags":["jax","ragged-dot","shape-mismatch","moe"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}