jax-ml/jax · error · ValueError

This gmm kernel only supports either (m, k) x (g, k, n) -> (

Error message

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=}

What it means

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.

Source

Thrown at jax/_src/lax/pallas_lowerings/gpu/ragged_dot.py:225

  group_sizes: Array,  # [g]
  block_m: int = DEFAULT_BLOCK_M,
  block_k: int = DEFAULT_BLOCK_K,
  block_n: int = DEFAULT_BLOCK_N,
  trans_rhs: bool = False,
  interpret: bool = False,
  compute_dtype: DTypeLike | None = None,
  acc_dtype: DTypeLike | None = np.float32,
  num_warps: int | None = None,
  num_stages: int | None = None,
  chunk_m: int = CHUNK_M,
  out_dtype: DTypeLike | None = None,
) -> Array:
  """Compute grouped matmul on GPU via a Pallas lowering."""

  msg = "This gmm kernel only supports either (m, k) x (g, k, n) -> (m, n) "
  msg += f"or (m, k) x (g, n, k) -> (m, n), but got {x.shape=} {A.shape=}"
  if not (A.ndim == 3 and x.ndim == 2):
    raise ValueError(msg)
  msg = f"Group sizes {group_sizes.shape=} must match first dimension of "
  msg += f"{A.shape=}"
  if not A.shape[:1] == group_sizes.shape:
    raise ValueError(msg)
  n = A.shape[-1] if not trans_rhs else A.shape[-2]
  Ak = A.shape[-2] if not trans_rhs else A.shape[-1]
  assert Ak == x.shape[1], msg
  size = RaggedDotSizes(m=x.shape[0], k=x.shape[1], n=n, g=A.shape[0])

  # normalize the block sizes for GPU
  block_m, block_k, block_n = (
    pl.next_power_of_2(min(b, s))
    for b, s in zip([block_m, block_k, block_n], [size.m, size.k, size.n])
  )
  block_k, block_n = max(block_k, 16), max(block_n, 16)

  A_spec = pl.BlockSpec((size.g, size.k, block_n), lambda i, j: (0, 0, j))
  if trans_rhs:  # transposed spec

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
out = gmm(x[None], A, group_sizes)  # x batched -> x.ndim==3

// after
out = gmm(x.reshape(x.shape[-2], x.shape[-1]), A, group_sizes)

Defensive patterns

Strategy: validation

Validate before calling

assert x.ndim == 2 and A.ndim == 3, f'gmm needs (m,k)x(g,k,n); got {x.shape}, {A.shape}'

Try / catch

try:
    out = gmm(x, A, group_sizes)
except ValueError as e:
    raise ValueError(f'reshape inputs for gmm: {e}') from e

Prevention

Keep activations 2-D and weights 3-D by convention in MoE code
Document the expected gmm shapes next to wrapper functions

When it happens

Trigger: 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).

Common situations: 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.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/1ff0d81210f52f4a. Report an issue: GitHub.