jax-ml/jax · error · NotImplementedError

Cannot partition grid over dynamic number of cores.

Error message

Cannot partition grid over dynamic number of cores.

What it means

_partition_grid splits a kernel grid across TPU cores for data-parallel execution. It requires num_cores to be a statically known Python int so it can divide grid dimensions; a tracer/Array value makes partitioning impossible, so it raises NotImplementedError.

Source

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

  elif isinstance(core_axis, int):
    return num_programs(core_axis), program_id(core_axis)
  else:
    return jax.lax.axis_size(core_axis), jax.lax.axis_index(core_axis)

def _partition_grid(
    grid: tuple[int | jax.Array, ...],
    dimension_semantics: tuple[GridDimensionSemantics, ...] | None,
    num_cores: int | None = None,
    core_id: jax.Array | int | None = None,
) -> tuple[tuple[int | jax.Array, ...], tuple[int | jax.Array, ...]]:
  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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert num_cores to a Python int before the call (e.g. int(num_cores) outside jit)
2. Mark num_cores/core_id as static arguments in any jit/partial wrapper
3. Or omit num_cores/core_id entirely to skip grid partitioning

Example fix

# before
num_cores = jax.numpy.size(devices)  # JAX scalar
emit_pipeline(..., num_cores=num_cores, core_id=core_id)
# after
num_cores = len(devices)  # Python int
emit_pipeline(..., num_cores=num_cores, core_id=core_id)

if num_cores is not None:
    assert isinstance(num_cores, int) and not hasattr(num_cores, 'aval'), 'num_cores must be a Python int'

def is_static_int(x) -> bool:
    return isinstance(x, int) and not isinstance(x, bool) or type(x).__name__ == 'int'

When it happens

Trigger: Passing num_cores (and core_id) derived from a JAX array or tracer (e.g. jax.numpy size, a value computed inside jit) rather than a Python int, together with dimension_semantics.

Common situations: Computing core counts dynamically from device topology arrays; passing core_id/num_cores through a jitted wrapper where they become tracers.

Related errors

• 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
• Acc ref must be at least 2D, got shape {shape}
• Acc ref dtype must be float32 or int32, got {dtype}

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