jax-ml/jax · error · ValueError

You can't use two different ScalarSubcoreMeshes.

Error message

You can't use two different ScalarSubcoreMeshes.

What it means

Raised by ScalarSubcoreMesh.check_is_compatible_with when a Pallas Mosaic SparseCore kernel's mesh is combined with another ScalarSubcoreMesh. The mesh compatibility API rejects using two distinct scalar subcore meshes in the same computation, because there is only one scalar subcore per SparseCore.

Source

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

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

  @property
  def size(self) -> int:
    return self.num_cores

  @property
  def dimension_semantics(self) -> Sequence[tpu_core.DimensionSemantics]:
    return [tpu_core.GridDimensionSemantics.CORE_PARALLEL]

  def discharges_effect(self, effect):
    del effect  # Unused.
    return False

  def check_is_compatible_with(self, other_mesh):
    if isinstance(other_mesh, ScalarSubcoreMesh):
      raise ValueError("You can't use two different ScalarSubcoreMeshes.")
    elif isinstance(other_mesh, VectorSubcoreMesh):
      if (self.axis_name == other_mesh.core_axis_name
          and self.num_cores == other_mesh.num_cores):
        return True
      raise ValueError(f"{self} should have the same core axis name and number"
                       f" of cores as the VectorSubcoreMesh {other_mesh}.")
    elif isinstance(other_mesh, tpu_core.TensorCoreMesh):
      if self.axis_name == other_mesh.axis_name:
        raise ValueError(
            f"{self} should have a different axis name from the TensorCoreMesh"
            f" {other_mesh}."
        )
      return True
    return super().check_is_compatible_with(other_mesh)

  @property
  def supported_memory_spaces(self) -> Sequence[Any]:
    return [

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a single shared ScalarSubcoreMesh instance for all participants
2. If you meant to express per-core parallelism, use VectorSubcoreMesh on the other side
3. Check the mesh types you pass to check_is_compatible_with before combining

Example fix

# before
mesh_a = pl_mosaic.ScalarSubcoreMesh(...)
mesh_b = pl_mosaic.ScalarSubcoreMesh(...)
mesh_a.check_is_compatible_with(mesh_b)  # raises
# after
mesh_a = pl_mosaic.ScalarSubcoreMesh(...)
mesh_a.check_is_compatible_with(vector_mesh)  # VectorSubcoreMesh with matching axis/cores

def is_scalar_mesh(m): return isinstance(m, pl_mosaic.ScalarSubcoreMesh)

import jax._src.pallas.mosaic.sc_core as sc
def is_scalar_subcore_mesh(m) -> bool:
    return isinstance(m, sc.ScalarSubcoreMesh)

try:
    a.check_is_compatible_with(b)
except ValueError as e:
    raise ConfigError(f'Mesh mismatch: {e}') from e

When it happens

Trigger: Calling check_is_compatible_with on a ScalarSubcoreMesh with another ScalarSubcoreMesh instance; e.g. combining SparseCore pallas grids/meshes where both sides are scalar subcore meshes.

Common situations: Building a multi-mesh Pallas TPU kernel (mixing TensorCore and SparseCore meshes) and accidentally passing two scalar subcore meshes; refactoring mesh setup and duplicating ScalarSubcoreMesh creation.

Related errors

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

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