jax-ml/jax · critical · RuntimeError

The current TPU does not have SparseCores

Error message

The current TPU does not have SparseCores

What it means

get_sparse_core_info queries the current TPU's SparseCore metadata; on chips without SparseCores (e.g. TPU v2/v3 or CPU/localhost) sparse_core is None and a RuntimeError is raised. Most SC-core primitives and mesh classes call this at construction time.

Source

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

import jax
from jax._src import core as jax_core
from jax._src import tree_util
from jax._src.pallas import core as pallas_core
from jax._src.pallas.mosaic import core as tpu_core
from jax._src.pallas.mosaic import tpu_info
import jax.numpy as jnp


def get_sparse_core_info() -> tpu_info.SparseCoreInfo:
  """Returns the SparseCore information for the current device.

  Raises:
    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"
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Run on TPU hardware that has SparseCores (v4/v5p/v6e etc.) by selecting jax.devices('tpu')
  2. Guard the mesh construction with get_sparse_core_info() in a try/except and skip SC paths when unavailable
  3. In tests/CI, gate SparseCore tests on device availability

Example fix

# before
mesh = VectorSubcoreMesh(...)
# after
try:
  info = get_sparse_core_info()
except RuntimeError:
  info = None
if info is not None:
  mesh = VectorSubcoreMesh(...)
Defensive patterns

Strategy: try-catch

Validate before calling

import jax
devs = jax.devices()
has_sc = any(d.platform == 'tpu' for d in devs)
# and expect tpu_info.sparse_core to be non-None on SC-capable chips

Type guard

def sparse_cores_available() -> bool:
    try:
        get_sparse_core_info()
        return True
    except RuntimeError:
        return False

Try / catch

try:
    sc_info = get_sparse_core_info()
except RuntimeError as e:
    if 'SparseCores' in str(e):
        raise SystemExit('This kernel requires a SparseCore-capable TPU') from e
    raise

Prevention

When it happens

Trigger: Constructing ScalarSubcoreMesh/VectorSubcoreMesh, calling supported_shapes, or using SC-core lowering rules on hardware without SparseCores, or calling tpu_info on a non-TPU host.

Common situations: Running a SparseCore Pallas kernel locally or in CI on CPU; targeting TPU v4+ only code on older TPU generations; using the wrong backend (gpu/cpu instead of tpu).

Related errors


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