{"record":{"id":"1ff0d81210f52f4a","repo":"jax-ml/jax","slug":"this-gmm-kernel-only-supports-either-m-k-x-g","errorCode":null,"errorMessage":"This gmm kernel only supports either (m, k) x (g, k, n) -> (m, n) or (m, k) x (g, n, k) -> (m, n), but got {x.shape=} {A.shape=}","messagePattern":"This gmm kernel only supports either \\(m, k\\) x \\(g, k, n\\) -> \\(m, n\\) or \\(m, k\\) x \\(g, n, k\\) -> \\(m, n\\), but got (.+?) (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/pallas_lowerings/gpu/ragged_dot.py","lineNumber":225,"sourceCode":"  group_sizes: Array,  # [g]\n  block_m: int = DEFAULT_BLOCK_M,\n  block_k: int = DEFAULT_BLOCK_K,\n  block_n: int = DEFAULT_BLOCK_N,\n  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","sourceCodeStart":207,"sourceCodeEnd":243,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/pallas_lowerings/gpu/ragged_dot.py#L207-L243","documentation":"The Pallas GPU grouped-matmul (gmm) kernel only accepts a 2-D activation matrix x of shape (m, k) multiplied by a 3-D group weights tensor A of shape (g, k, n) or (g, n, k). Any other rank combination (e.g. batched x with ndim 3, or 2-D weights) raises this ValueError before kernel launch.","triggerScenarios":"Calling jax.experimental.pallas...gmm / the Mosaic GPU gmm lowering with x.ndim != 2 or A.ndim != 3, e.g. passing a batched (b, m, k) activation or a single (k, n) weight matrix instead of (g, k, n).","commonSituations":"Porting a batched dot to MoI/MoE grouped matmul without reshaping; passing trans_rhs with weights stored 2-D; using gmm where a vmap over a plain matmul was used before.","solutions":["Reshape activations to 2-D: x.reshape(-1, k) before gmm","Add a leading group dimension to the weights so A has shape (g, k, n) (or (g, n, k) with trans_rhs=True)","If you truly need per-batch matmuls, use lax.dot_general / jnp.einsum instead of gmm"],"exampleFix":"// before\nout = gmm(x[None], A, group_sizes)  # x batched -> x.ndim==3\n\n// after\nout = gmm(x.reshape(x.shape[-2], x.shape[-1]), A, group_sizes)","handlingStrategy":"validation","validationCode":"assert x.ndim == 2 and A.ndim == 3, f'gmm needs (m,k)x(g,k,n); got {x.shape}, {A.shape}'","typeGuard":null,"tryCatchPattern":"try:\n    out = gmm(x, A, group_sizes)\nexcept ValueError as e:\n    raise ValueError(f'reshape inputs for gmm: {e}') from e","preventionTips":["Keep activations 2-D and weights 3-D by convention in MoE code","Document the expected gmm shapes next to wrapper functions"],"tags":["jax","pallas","gpu","shape","moe"],"backgroundTag":"invalid-shape-for-kernel","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}