{"record":{"id":"3e62df9484eeb6b2","repo":"jax-ml/jax","slug":"cannot-partition-grid-over-dynamic-number-of-cores","errorCode":null,"errorMessage":"Cannot partition grid over dynamic number of cores.","messagePattern":"Cannot partition grid over dynamic number of cores\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/pipeline.py","lineNumber":1553,"sourceCode":"  elif isinstance(core_axis, int):\n    return num_programs(core_axis), program_id(core_axis)\n  else:\n    return jax.lax.axis_size(core_axis), jax.lax.axis_index(core_axis)\n\ndef _partition_grid(\n    grid: tuple[int | jax.Array, ...],\n    dimension_semantics: tuple[GridDimensionSemantics, ...] | None,\n    num_cores: int | None = None,\n    core_id: jax.Array | int | None = None,\n) -> tuple[tuple[int | jax.Array, ...], tuple[int | jax.Array, ...]]:\n  assert not ((num_cores is None) ^ (core_id is None)), (\n      \"Either both num_cores and core_id should be provided, or neither.\")\n  if num_cores is None or core_id is None:\n    # We aren't partitioning the grid\n    return grid, (0,) * len(grid)\n  # Check that num_cores is statically known\n  if not isinstance(num_cores, int):\n    raise NotImplementedError(\n        \"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","sourceCodeStart":1535,"sourceCodeEnd":1571,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/pipeline.py#L1535-L1571","documentation":"_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.","triggerScenarios":"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.","commonSituations":"Computing core counts dynamically from device topology arrays; passing core_id/num_cores through a jitted wrapper where they become tracers.","solutions":["Convert num_cores to a Python int before the call (e.g. int(num_cores) outside jit)","Mark num_cores/core_id as static arguments in any jit/partial wrapper","Or omit num_cores/core_id entirely to skip grid partitioning"],"exampleFix":"# before\nnum_cores = jax.numpy.size(devices)  # JAX scalar\nemit_pipeline(..., num_cores=num_cores, core_id=core_id)\n# after\nnum_cores = len(devices)  # Python int\nemit_pipeline(..., num_cores=num_cores, core_id=core_id)","handlingStrategy":"type-guard","validationCode":"if num_cores is not None:\n    assert isinstance(num_cores, int) and not hasattr(num_cores, 'aval'), 'num_cores must be a Python int'","typeGuard":"def is_static_int(x) -> bool:\n    return isinstance(x, int) and not isinstance(x, bool) or type(x).__name__ == 'int'","tryCatchPattern":null,"preventionTips":["Convert topology-derived core counts with int(...) outside any jit","Mark num_cores/core_id as static in jit wrappers"],"tags":["jax","pallas","tpu","grid-partition","static-value"],"backgroundTag":"dynamic-value-where-static-required","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}