jax-ml/jax · error · ValueError

Dimensions with parallel semantics must form a prefix of the

Error message

Dimensions with parallel semantics must form a prefix of the grid.

What it means

For randomized TPU interpret mode (multi-core simulation), the grid's dimension_semantics must have all 'parallel' dimensions as a prefix — no 'parallel' dimension may follow a non-parallel one. Otherwise the interpreter cannot compute consistent randomized grid coordinates and raises this ValueError.

Source

Thrown at jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py:1822

  Args:
    mosaic_params: The compiler params for the Mosaic TPU backend.
    num_dimensions_in_grid: The number of dimensions in the grid.

  Returns:
    A tuple of booleans where the entry at index `i` is `True` precisely if the
    `i`-th dimension in the grid has parallel semantics.

  Raises:
    ValueError: If the dimensions with parallel semantics do not form a prefix
      of the grid.
  """
  if mosaic_params.dimension_semantics is None:
    return (False,) * num_dimensions_in_grid
  result = tuple(ds in ('parallel', mosaic_core.PARALLEL)
                 for ds in mosaic_params.dimension_semantics)
  for ds0, ds1 in zip(result[:-1], result[1:]):
    if ds1 and not ds0:
      raise ValueError(
          'Dimensions with parallel semantics must form a prefix of the grid.'
      )
  return result


def _get_parallel_subgrid_size(
    parallel_semantics_per_dim: tuple[bool, ...], grid: tuple[int, ...]
) -> int:
  """Returns the size of the subgrid along the parallel dimensions."""
  return math.prod(
      dim_size if parallel_dim else 1
      for dim_size, parallel_dim in zip(grid, parallel_semantics_per_dim)
  )

_GridPointCoordinatesPerDim = tuple[Array, ...]

def _get_randomized_grid_coordinates(
    grid: tuple[int, ...],

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reorder the grid so parallel dimensions come first and relabel dimension_semantics accordingly
2. Mark the trailing dimensions as 'loop' if they carry sequential dependencies
3. Permute grid axes (and corresponding BlockSpec index_maps) to satisfy the prefix rule

Example fix

# before
dimension_semantics=['parallel', 'loop', 'parallel']
# after
dimension_semantics=['parallel', 'parallel', 'loop']  # reorder grid accordingly

def semantics_ok(ds):
    seen_non_parallel = False
    for d in ds:
        if d != 'parallel':
            seen_non_parallel = True
        elif seen_non_parallel:
            return False
    return True
assert semantics_ok(dimension_semantics)

def is_parallel_prefix(ds) -> bool:
    p = [d in ('parallel',) for d in ds]
    return all(p[:sum(p)]) and not any(p[sum(p):])

try:
    interpret_run(kernel)
except ValueError as e:
    if 'prefix of the grid' in str(e):
        # reorder grid dims so parallel dims come first, update index_maps
        raise

When it happens

Trigger: Calling interpret_pallas_call with CompilerParams(dimension_semantics=['parallel','loop','parallel']) or any layout where a parallel dim comes after a non-parallel dim.

Common situations: Matmul-style kernels with contraction dims marked loop followed by parallel broadcast dims; reordering grid dims for performance while keeping old semantics labels; multi-core interpret runs of collective kernels.

Related errors

• Revisited block {output_ranges[i]} of output {i} in iteratio
• Non-trivial windowing is not supported for grid-free pallas_
• Cannot partition over cores without parallel grid dimensions
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s

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