jax-ml/jax · error · NotImplementedError

Cannot partition over cores without parallel grid dimensions

Error message

Cannot partition over cores without parallel grid dimensions: {dimension_semantics=}

What it means

Grid partitioning across cores can only split axes marked PARALLEL in dimension_semantics. If no axis is parallel, there is nothing to shard, and since running on a single core is not yet supported (per the TODO), this NotImplementedError is raised.

Source

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

        "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
  }
  if divisible_dimensions:
    first_divisible_dimension, *_ = (
        i for i in range(len(dimension_semantics)) if i in divisible_dimensions
    )
    partitioned_dim_size = grid[first_divisible_dimension] // num_cores
    partitioned_dim_offset = core_id * partitioned_dim_size
    new_grid = jax_util.tuple_update(
        grid, first_divisible_dimension, partitioned_dim_size

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Mark the data-parallel grid axis as 'parallel' in dimension_semantics
2. Ensure every axis you want sharded is labeled PARALLEL and has divisible size
3. If single-core execution is fine, drop num_cores/core_id arguments

Example fix

# before
emit_pipeline(..., num_cores=4, core_id=0)  # no dimension_semantics
# after
emit_pipeline(..., num_cores=4, core_id=0,
              dimension_semantics=('parallel', 'arbitrary'))

if num_cores is not None:
    sems = dimension_semantics or ('arbitrary',) * len(grid)
    assert any(d == 'parallel' for d in sems), 'partitioning needs >=1 parallel axis'

When it happens

Trigger: Providing num_cores/core_id with dimension_semantics containing only ARBITRARY entries, or with dimension_semantics=None (defaults to all ARBITRARY).

Common situations: Forgetting to pass dimension_semantics when enabling multi-core partitioning; assuming arbitrary axes can be sharded.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• Dimensions with parallel semantics must form a prefix of the
• Cannot partition grid over dynamic number of cores.
• 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/e5d4b965cc3d063a. Report an issue: GitHub.