jax-ml/jax · error · ValueError

No tuned tiling found for (m, k, n) = ({m}, {k}, {n})

Error message

No tuned tiling found for (m, k, n) = ({m}, {k}, {n})

What it means

Megablox gmm selects tile sizes from a tuned lookup table (LUT) keyed by (m, k, n) problem dimensions. When the tiling argument is a callable LUT and it returns None for the given dimensions, no tuned configuration exists and the kernel refuses to guess, raising this ValueError.

Source

Thrown at jax/experimental/pallas/ops/tpu/megablox/gmm.py:376

    group_offset = group_offset[None]
  num_current_groups = rhs.shape[0]
  num_total_groups = group_sizes.shape[0]
  lhs, group_sizes, input_dtype = _validate_args(
      lhs=lhs, rhs=rhs, group_sizes=group_sizes
  )

  # Gather shape information.
  m, k, n = (lhs.shape[0], lhs.shape[1], rhs.shape[2])
  if transpose_rhs:
    n = rhs.shape[1]

  # If tiling is callable, look up the problem dimensions in the LUT. If no tuned
  # tile dimensions are available throw an error.
  if callable(tiling):
    tiling = tiling(m, k, n)

  if tiling is None:
    raise ValueError(f"No tuned tiling found for (m, k, n) = ({m}, {k}, {n})")

  tm, tk, tn = tiling
  tiles_k, k_rem = _calculate_irregular_num_tiles(k, tk)
  tiles_n, n_rem = _calculate_irregular_num_tiles(n, tn)
  del n_rem

  # Create the metadata we need for computation.
  group_metadata, num_active_tiles = make_group_metadata(
      group_sizes=group_sizes,
      m=m,
      tm=tm,
      start_group=group_offset[0],
      num_nonzero_groups=rhs.shape[0],
      visit_empty_groups=False,
  )

  def kernel(
      group_metadata,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass an explicit tiling tuple, e.g. tiling=(128, 128, 128), instead of the LUT callable
2. Round m/k/n to dimensions present in the tuned table (typically multiples of 128)
3. Use jnp.einsum/dot as a fallback for untuned shapes

Example fix

// before
out = gmm(lhs, rhs, group_sizes)  # default LUT
// after
out = gmm(lhs, rhs, group_sizes, tiling=(128, 128, 128))

tiling = select_tiling(m, k, n)
if tiling is None:
    tiling = (128, 128, 128)  # safe default for your workload

try:
    out = gmm(lhs, rhs, group_sizes)
except ValueError as e:
    if 'No tuned tiling' in str(e):
        out = gmm(lhs, rhs, group_sizes, tiling=(128, 128, 128))
    else:
        raise

When it happens

Trigger: Calling gmm with default tiling=select_tiling (or another callable) for a matrix shape absent from the tuning LUT, e.g. unusual m=1333 or k dimensions not covered by tuning runs. Also raised by _gmm_fwd/_gmm_bwd during autodiff of such a call.

Common situations: Changing batch size or model dims to a value never tuned (e.g. non-multiples of 128); using a newer/older JAX whose LUT covers a different shape set; running MoE experts with atypical shapes.

Related errors

• Existing output dtype must match preferred_element_type.
• group_offset must be a ()-shaped array. Got: {group_offset.s
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s
• Acc ref must be at least 2D, got shape {shape}

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