jax-ml/jax · error · ValueError

dimension_semantics must be the same length as grid.

Error message

dimension_semantics must be the same length as grid.

What it means

dimension_semantics labels each grid axis as PARALLEL or ARBITRARY and must have exactly one entry per grid dimension. If the tuple length differs from len(grid), partitioning cannot align semantics with axes and this ValueError is raised.

Source

Thrown at jax/_src/pallas/mosaic/pipeline.py:1564

  assert not ((num_cores is None) ^ (core_id is None)), (
      "Either both num_cores and core_id should be provided, or neither.")
  if num_cores is None or core_id is None:
    # We aren't partitioning the grid
    return grid, (0,) * len(grid)
  # Check that num_cores is statically known
  if not isinstance(num_cores, int):
    raise NotImplementedError(
        "Cannot partition grid over dynamic number of cores."
    )
  if num_cores == 1:
    # We aren't partitioning the grid
    return grid, (0,) * len(grid)

  # If dimension_semantics aren't provided, we assume it is all arbitrary.
  if dimension_semantics is None:
    dimension_semantics = (ARBITRARY,) * len(grid)
  if len(dimension_semantics) != len(grid):
    raise ValueError("dimension_semantics must be the same length as grid.")

  parallel_dimensions = {
      i for i, d in enumerate(dimension_semantics) if d == PARALLEL
  }
  # If there are no parallel dimensions, we can't partition the grid
  if not parallel_dimensions:
    # TODO(sharadmv): enable running kernel on just one core
    raise NotImplementedError(
        "Cannot partition over cores without parallel grid dimensions:"
        f" {dimension_semantics=}"
    )

  # Try to find a divisible dimension to partition the grid on
  divisible_dimensions = {
      i
      for i in parallel_dimensions
      if isinstance(grid[i], int) and grid[i] % num_cores == 0
  }

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make dimension_semantics length match len(grid), one label per axis
2. Omit dimension_semantics to get the all-ARBITRARY default (note: partitioning then requires no parallel dims, see error 2228)
3. Build it programmatically: ('parallel',) * len(grid) when all axes are parallel

Example fix

# before
grid=(4, 8, 2), dimension_semantics=('parallel', 'arbitrary')
# after
grid=(4, 8, 2), dimension_semantics=('parallel', 'arbitrary', 'parallel')

assert dimension_semantics is None or len(dimension_semantics) == len(grid), \
    f'{len(dimension_semantics)} semantics vs grid of rank {len(grid)}'

When it happens

Trigger: Calling the pipeline API with dimension_semantics shorter/longer than grid, e.g. grid=(4, 8, 2) with dimension_semantics=('parallel', 'arbitrary').

Common situations: Adding a grid dimension (e.g. batching) without updating dimension_semantics; defaulting semantics for a different grid shape than the one passed.

Related errors

• Sum of sizes {n} must be equal to dimension {axis} of the op
• Cannot pull iota along dimension {dimension} with None block
• Invalid memory space: {memory_space!r}
• Dimensions with parallel semantics must form a prefix of the
• Cannot partition over cores without parallel grid dimensions

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