jax-ml/jax · error · NotImplementedError

Cannot partition over multiple dynamic parallel dimensions

Error message

Cannot partition over multiple dynamic parallel dimensions: {grid=}

What it means

When partitioning a grid with a dynamic (non-Python-int) dimension, the partitioner can handle at most one dynamic parallel dimension. If two or more parallel axes are dynamic JAX arrays, it cannot statically distribute work, so NotImplementedError is raised.

Solutions

1. Make all but at most one parallel grid dimension a static Python int
2. Convert dynamic sizes with int(...) where the value is known at trace time
3. Disable multi-core partitioning (omit num_cores/core_id) for fully dynamic grids

Example fix

# before
grid=(seq_len_jax, batch_jax, 8)  # two dynamic parallel dims
# after
grid=(int(seq_len), batch_jax, 8)  # one dynamic dim max

sem = dimension_semantics or ()
dyn = [i for i, (g, d) in enumerate(zip(grid, sem)) if d == 'parallel' and not isinstance(g, int)]
assert len(dyn) <= 1, f'multiple dynamic parallel dims: {dyn}'

def at_most_one_dynamic_parallel(grid, sems) -> bool:
    return sum(1 for g, d in zip(grid, sems) if d == 'parallel' and not isinstance(g, int)) <= 1

When it happens

Trigger: Passing a grid like (jax_arr_size, jax_arr_size2, 8) where two of the parallel-marked dimensions are JAX Arrays rather than Python ints, together with num_cores/core_id.

Common situations: Fully dynamic grids derived from runtime shapes in multi-core TPU kernels; migrating a static-grid kernel to dynamic sizing while keeping core partitioning on.

Related errors

• Cannot partition grid over dynamic number of cores.
• Cannot partition over cores without parallel grid dimensions
• The iteration bound corresponding to the core-parallel…
• a and b must be 2D, but got
• a and b must have the same element type, but got

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

Appendix: source

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

        first_divisible_dimension,
        partitioned_dim_offset,
    )
    return new_grid, offsets

  # Separate the remaining dimensions into dynamic and static.
  dynamic_dims = [
      i
      for i in range(len(grid))
      if i in parallel_dimensions and not isinstance(grid[i], int)
  ]
  static_dims = [
      i
      for i in range(len(grid))
      if i in parallel_dimensions and isinstance(grid[i], int)
  ]

  if len(dynamic_dims) > 1:
    raise NotImplementedError(
        f"Cannot partition over multiple dynamic parallel dimensions: {grid=}"
    )

  if dynamic_dims and not static_dims:
    # Exactly one dynamic dimension and no static non-divisible dimensions
    partition_dimension = dynamic_dims[0]
  else:
    # No divisible static dimensions, so we can't evenly partition the grid.
    # Let's pick the largest dimension and try to divide it as evenly as
    # possible.
    # TODO(sharadmv): take the product of many nondivisible dimensions to
    # potentially divide it more evenly
    largest_parallel_dimension = max(grid[i] for i in static_dims)
    partition_dimension, *_ = (
        i for i in static_dims if grid[i] == largest_parallel_dimension
    )

  base_num_iters, rem = divmod(grid[partition_dimension], num_cores)

View on GitHub (pinned to 1e1c6a8fc0)