{"record":{"id":"de5c384716708b05","repo":"jax-ml/jax","slug":"group-sizes-group-sizes-shape-must-match-first","errorCode":null,"errorMessage":"Group sizes {group_sizes.shape=} must match first dimension of {A.shape=}","messagePattern":"Group sizes (.+?) must match first dimension of (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/pallas_lowerings/gpu/ragged_dot.py","lineNumber":229,"sourceCode":"  trans_rhs: bool = False,\n  interpret: bool = False,\n  compute_dtype: DTypeLike | None = None,\n  acc_dtype: DTypeLike | None = np.float32,\n  num_warps: int | None = None,\n  num_stages: int | None = None,\n  chunk_m: int = CHUNK_M,\n  out_dtype: DTypeLike | None = None,\n) -> Array:\n  \"\"\"Compute grouped matmul on GPU via a Pallas lowering.\"\"\"\n\n  msg = \"This gmm kernel only supports either (m, k) x (g, k, n) -> (m, n) \"\n  msg += f\"or (m, k) x (g, n, k) -> (m, n), but got {x.shape=} {A.shape=}\"\n  if not (A.ndim == 3 and x.ndim == 2):\n    raise ValueError(msg)\n  msg = f\"Group sizes {group_sizes.shape=} must match first dimension of \"\n  msg += f\"{A.shape=}\"\n  if not A.shape[:1] == group_sizes.shape:\n    raise ValueError(msg)\n  n = A.shape[-1] if not trans_rhs else A.shape[-2]\n  Ak = A.shape[-2] if not trans_rhs else A.shape[-1]\n  assert Ak == x.shape[1], msg\n  size = RaggedDotSizes(m=x.shape[0], k=x.shape[1], n=n, g=A.shape[0])\n\n  # normalize the block sizes for GPU\n  block_m, block_k, block_n = (\n    pl.next_power_of_2(min(b, s))\n    for b, s in zip([block_m, block_k, block_n], [size.m, size.k, size.n])\n  )\n  block_k, block_n = max(block_k, 16), max(block_n, 16)\n\n  A_spec = pl.BlockSpec((size.g, size.k, block_n), lambda i, j: (0, 0, j))\n  if trans_rhs:  # transposed spec\n    A_spec = pl.BlockSpec((size.g, block_n, size.k), lambda i, j: (0, j, 0))\n\n  group_metadata = _make_gmm_group_metadata(\n    group_sizes=group_sizes, m=size.m, chunk_m=chunk_m","sourceCodeStart":211,"sourceCodeEnd":247,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/pallas_lowerings/gpu/ragged_dot.py#L211-L247","documentation":"In the Pallas GPU grouped-matmul kernel, group_sizes must be a 1-D array whose length equals the number of groups g (the first dimension of the 3-D weights A). A mismatch means the kernel cannot map each group to its row-range in the ragged output.","triggerScenarios":"Calling gmm with A.shape == (g, k, n) but group_sizes.shape != (g,), e.g. passing per-expert sizes of length num_experts while A packs a different number of groups (or passing a scalar/split array).","commonSituations":"MoE routing where group_sizes comes from a routing histogram over a different expert count than the weight tensor; off-by-one or including a padding group in A but not in group_sizes.","solutions":["Verify A.shape[0] == group_sizes.shape[0] and reconcile the expert count on both sides","Build group_sizes from the same routing computation that permuted the weights (e.g. cumsum-derived sizes of the sorted token assignment)","Add an explicit assert before gmm to fail early with your own message"],"exampleFix":"// before\nout = gmm(x, A, group_sizes)  # len(group_sizes) != A.shape[0]\n\n// after\nassert group_sizes.shape[0] == A.shape[0]\nout = gmm(x, A, group_sizes)","handlingStrategy":"validation","validationCode":"assert group_sizes.shape == A.shape[:1], (group_sizes.shape, A.shape)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive group_sizes and A from the same routing state","Add shape asserts in a thin gmm wrapper used everywhere"],"tags":["jax","pallas","moe","shape","gpu"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}