{"record":{"id":"e5d4b965cc3d063a","repo":"jax-ml/jax","slug":"cannot-partition-over-cores-without-parallel-grid","errorCode":null,"errorMessage":"Cannot partition over cores without parallel grid dimensions: {dimension_semantics=}","messagePattern":"Cannot partition over cores without parallel grid dimensions: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/pipeline.py","lineNumber":1572,"sourceCode":"        \"Cannot partition grid over dynamic number of cores.\"\n    )\n  if num_cores == 1:\n    # We aren't partitioning the grid\n    return grid, (0,) * len(grid)\n\n  # If dimension_semantics aren't provided, we assume it is all arbitrary.\n  if dimension_semantics is None:\n    dimension_semantics = (ARBITRARY,) * len(grid)\n  if len(dimension_semantics) != len(grid):\n    raise ValueError(\"dimension_semantics must be the same length as grid.\")\n\n  parallel_dimensions = {\n      i for i, d in enumerate(dimension_semantics) if d == PARALLEL\n  }\n  # If there are no parallel dimensions, we can't partition the grid\n  if not parallel_dimensions:\n    # TODO(sharadmv): enable running kernel on just one core\n    raise NotImplementedError(\n        \"Cannot partition over cores without parallel grid dimensions:\"\n        f\" {dimension_semantics=}\"\n    )\n\n  # Try to find a divisible dimension to partition the grid on\n  divisible_dimensions = {\n      i\n      for i in parallel_dimensions\n      if isinstance(grid[i], int) and grid[i] % num_cores == 0\n  }\n  if divisible_dimensions:\n    first_divisible_dimension, *_ = (\n        i for i in range(len(dimension_semantics)) if i in divisible_dimensions\n    )\n    partitioned_dim_size = grid[first_divisible_dimension] // num_cores\n    partitioned_dim_offset = core_id * partitioned_dim_size\n    new_grid = jax_util.tuple_update(\n        grid, first_divisible_dimension, partitioned_dim_size","sourceCodeStart":1554,"sourceCodeEnd":1590,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/pipeline.py#L1554-L1590","documentation":"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.","triggerScenarios":"Providing num_cores/core_id with dimension_semantics containing only ARBITRARY entries, or with dimension_semantics=None (defaults to all ARBITRARY).","commonSituations":"Forgetting to pass dimension_semantics when enabling multi-core partitioning; assuming arbitrary axes can be sharded.","solutions":["Mark the data-parallel grid axis as 'parallel' in dimension_semantics","Ensure every axis you want sharded is labeled PARALLEL and has divisible size","If single-core execution is fine, drop num_cores/core_id arguments"],"exampleFix":"# before\nemit_pipeline(..., num_cores=4, core_id=0)  # no dimension_semantics\n# after\nemit_pipeline(..., num_cores=4, core_id=0,\n              dimension_semantics=('parallel', 'arbitrary'))","handlingStrategy":"validation","validationCode":"if num_cores is not None:\n    sems = dimension_semantics or ('arbitrary',) * len(grid)\n    assert any(d == 'parallel' for d in sems), 'partitioning needs >=1 parallel axis'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Label data-parallel axes as 'parallel' before enabling num_cores/core_id","Remember the all-arbitrary default cannot be partitioned"],"tags":["jax","pallas","tpu","grid-partition","dimension-semantics"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}