jax-ml/jax · error · ValueError

m_warps must be 1, 2, or 4, but got {m_warps=}

Error message

m_warps must be 1, 2, or 4, but got {m_warps=}

What it means

FragmentedArray MMA layout construction partitions 4 warps between M and N dimensions; m_warps must therefore divide into 4, and only 1, 2, or 4 are valid. Other values raise ValueError.

Source

Thrown at jax/experimental/mosaic/gpu/mma.py:37

from jax.experimental.mosaic.gpu import fragmented_array as fa
from jaxlib.mlir import ir
from jaxlib.mlir.dialects import llvm
from jaxlib.mlir.dialects import vector
import numpy as np
from . import utils


SUPPORTED_F8_TYPES = (ir.Float8E4M3FNType, ir.Float8E5M2Type)


class MMALayouts:
  """Container for MMA layouts, providing a convenient way to create
  layouts for MMA operands based on warp configuration.
  """

  def __init__(self, element_type: ir.Type | typing.DTypeLike, *, m_warps: int = 4):
    if m_warps not in (1, 2, 4):
      raise ValueError(f"m_warps must be 1, 2, or 4, but got {m_warps=}")
    n_warps = 4 // m_warps
    if isinstance(element_type, ir.Type):
      bitwidth = utils.bitwidth(element_type)
    else:
      bitwidth = dtypes.itemsize_bits(element_type)
    elems_per_reg = 32 // bitwidth
    k = 8 * elems_per_reg
    sub_k = 4 * elems_per_reg
    self.lhs = fa.TiledLayout(
        fa.Tiling(((m_warps * 16, k), (16, sub_k), (8, sub_k), (elems_per_reg,))),
        warp_dims=(-7, fa.Replicated(4 // m_warps)),
        lane_dims=(-3, -2),
        vector_dim=-1,
        _check_canonical=False,
    ).canonicalize()
    self.rhs = fa.TiledLayout(
        fa.Tiling(((k, n_warps * 8), (sub_k, 8), (elems_per_reg, 1))),
        warp_dims=(fa.Replicated(4 // n_warps), -5,),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use m_warps in (1, 2, 4)
2. Validate user-supplied warp configs before passing them in
3. If you need more warps, scale via repetition/other tiling, not m_warps

Example fix

// before
layout = mma.MMALayout(dt, m_warps=3)
// after
layout = mma.MMALayout(dt, m_warps=2)

assert m_warps in (1, 2, 4), 'm_warps must divide 4'

When it happens

Trigger: Passing m_warps=3, 8, 0, etc. to an MMA layout constructor / FragmentedArray helper that takes m_warps.

Common situations: Parameterizing kernel warp splits from config strings or sweep scripts without validating against {1,2,4}.

Related errors

• is_signed must be specified for integer types
• M mismatch: {m} != {m2}
• N mismatch: {n} != {n2}
• K mismatch: {k} != {k2}
• Dtype mismatch: {a.mlir_dtype} != {b.mlir_dtype}

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