{"record":{"id":"56cd041c2e684128","repo":"jax-ml/jax","slug":"expected-group-sizes-to-have-shape-expected-gs-sh","errorCode":null,"errorMessage":"expected group_sizes to have shape {expected_gs_shape}, got {group_sizes.shape}.","messagePattern":"expected group_sizes to have shape (.+?), got (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":6400,"sourceCode":"  mode = _ragged_dot_mode(lhs.ndim, ragged_dot_dimension_numbers)\n\n  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = (\n      ragged_dot_dimension_numbers.dot_dimension_numbers\n  )\n\n  # Validate the shape of group_sizes, if it is something other than [g].\n  if group_sizes.ndim == 0:\n    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:","sourceCodeStart":6382,"sourceCodeEnd":6418,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L6382-L6418","documentation":"When group_sizes has rank > 1, ragged_dot_general requires its shape to be exactly the broadcast prefix (batch and non-contracting dims of lhs) plus the trailing group axis, i.e. [b..., x..., g]; anything else fails this check.","triggerScenarios":"Passing group_sizes shaped (g, b) instead of (b, g), or omitting lhs batch dims from group_sizes, in a batched ragged_dot_general call.","commonSituations":"Nested/moe-style grouped GEMMs where per-batch group sizes were stacked in the wrong axis order; forgetting to repeat group sizes along a non-contracting lhs axis.","solutions":["Transpose/reshape group_sizes so the group axis is last and lhs batch/non-contracting dims come first in lhs order","Broadcast a base (b..., 1) group-size array with jnp.broadcast_to to the expected (b..., g) shape"],"exampleFix":"# before\nout = lax.ragged_dot_general(x, y, gs, dims)  # gs.shape == (g, b)\n# after\ngs = gs.T  # (b, g)\nout = lax.ragged_dot_general(x, y, gs, dims)","handlingStrategy":"validation","validationCode":"if group_sizes.ndim > 1:\n    prefix = tuple(lhs.shape[i] for i in prefix_dims)\n    expected = prefix + (group_sizes.shape[-1],)\n    assert group_sizes.shape == expected, (group_sizes.shape, expected)\n# or broadcast defensively:\n# group_sizes = jnp.broadcast_to(group_sizes, expected)","typeGuard":"def group_sizes_shape_ok(lhs, gs, prefix_dims):\n    if gs.ndim <= 1: return True\n    return gs.shape == tuple(lhs.shape[i] for i in prefix_dims) + (gs.shape[-1],)","tryCatchPattern":null,"preventionTips":["Keep the group axis last in group_sizes","Use jnp.broadcast_to from (b..., 1) to (b..., g) for repeated sizes"],"tags":["jax","ragged-dot-general","group-sizes","shape-mismatch"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}