jax-ml/jax · error · ValueError

No tuned tiling found for (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.

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
• Acc ref dtype must be float32 or int32, got
• Acc ref must be at least 2D, got shape
• Accumulators are not available on TPU

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

Appendix: 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)