jax-ml/jax · error · ValueError

Mesh has {self.num_cores} cores, but the current TPU chip ha

Error message

Mesh has {self.num_cores} cores, but the current TPU chip has only {sc_info.num_cores} SparseCores

What it means

SparseCore mesh dataclasses (ScalarSubcoreMesh etc.) validate their num_cores against the actual SparseCore count of the current chip in __post_init__. Requesting more SC cores than physically present raises ValueError immediately.

Source

Thrown at jax/_src/pallas/mosaic/sc_core.py:55

    RuntimeError: If the current TPU does not have SparseCores.
  """
  sc_info = tpu_info.get_tpu_info().sparse_core
  if sc_info is None:
    raise RuntimeError("The current TPU does not have SparseCores")
  return sc_info


@dataclasses.dataclass(frozen=True, kw_only=True)
class ScalarSubcoreMesh(pallas_core.Mesh):
  axis_name: str
  num_cores: int = dataclasses.field(
      default_factory=lambda: get_sparse_core_info().num_cores
  )

  def __post_init__(self):
    sc_info = get_sparse_core_info()
    if self.num_cores > sc_info.num_cores:
      raise ValueError(
          f"Mesh has {self.num_cores} cores, but the current TPU chip has only"
          f" {sc_info.num_cores} SparseCores"
      )

  @property
  def core_type(self) -> tpu_core.CoreType:
    return tpu_core.CoreType.SC_SCALAR_SUBCORE

  @property
  def default_memory_space(self) -> tpu_core.MemorySpace:
    return tpu_core.MemorySpace.HBM

  @property
  def shape(self):
    return collections.OrderedDict({self.axis_name: self.num_cores})

  @property
  def size(self) -> int:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Omit num_cores to use the chip default: num_cores defaults to get_sparse_core_info().num_cores
  2. Set num_cores = min(requested, get_sparse_core_info().num_cores) computed at runtime

Example fix

# before
mesh = ScalarSubcoreMesh(axis_name='sc', num_cores=8)  # chip has 4
# after
mesh = ScalarSubcoreMesh(axis_name='sc')  # uses chip's SparseCore count
Defensive patterns

Strategy: validation

Validate before calling

sc_info = get_sparse_core_info()
num_cores = min(num_cores, sc_info.num_cores)

Try / catch

try:
    mesh = ScalarSubcoreMesh(axis_name='sc', num_cores=n)
except ValueError:
    mesh = ScalarSubcoreMesh(axis_name='sc')  # chip default

Prevention

When it happens

Trigger: Constructing a SparseCore mesh with num_cores greater than get_sparse_core_info().num_cores, e.g. num_cores=8 on a chip with 4 SparseCores.

Common situations: Hardcoding core counts from a different TPU generation; scaling num_cores with a config value tuned for a larger pod slice; defaulting num_cores from a stale cached tpu_info.

Related errors


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