jax-ml/jax · error · ValueError

Mesh has {self.num_subcores} subcores, but the current TPU c

Error message

Mesh has {self.num_subcores} subcores, but the current TPU chip has only {sc_info.num_subcores} subcores

What it means

VectorSubcoreMesh.__post_init__ validates num_subcores against the number of subcores per SparseCore on the current chip. A value larger than the hardware's subcore count fails at construction time.

Source

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

class VectorSubcoreMesh(pallas_core.Mesh):
  core_axis_name: str
  subcore_axis_name: str
  num_cores: int = dataclasses.field(
      default_factory=lambda: get_sparse_core_info().num_cores
  )
  num_subcores: int = dataclasses.field(
      default_factory=lambda: get_sparse_core_info().num_subcores
  )

  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"
      )
    if self.num_subcores > sc_info.num_subcores:
      raise ValueError(
          f"Mesh has {self.num_subcores} subcores, but the current TPU chip has"
          f" only {sc_info.num_subcores} subcores"
      )

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

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

  @property
  def shape(self):
    return collections.OrderedDict({
        self.core_axis_name: self.num_cores,
        self.subcore_axis_name: self.num_subcores,
    })

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Omit num_subcores so it defaults to the hardware value
2. Set num_subcores = get_sparse_core_info().num_subcores

Example fix

# before
mesh = pl_mosaic.VectorSubcoreMesh(..., num_subcores=8)
# after
mesh = pl_mosaic.VectorSubcoreMesh(..., num_subcores=get_sparse_core_info().num_subcores)

n = get_sparse_core_info().num_subcores
assert kwargs.get('num_subcores', n) <= n

null

null

When it happens

Trigger: Building a VectorSubcoreMesh whose num_subcores exceeds get_sparse_core_info().num_subcores; default (num_subcores default_factory) normally matches hardware, so this happens when overriding it.

Common situations: Manually setting num_subcores for a different chip generation; assuming 2x subcore count relative to actual hardware.

Related errors

• The current TPU does not have SparseCores
• Mesh has {self.num_cores} cores, but the current TPU chip ha
• You can't use two different ScalarSubcoreMeshes.
• {self} should have the same core axis name and number of cor
• {self} should have a different axis name from the TensorCore

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