jax-ml/jax · error · ValueError

{self} should have a different axis name from the TensorCore

Error message

{self} should have a different axis name from the TensorCoreMesh {other_mesh}.

What it means

A ScalarSubcoreMesh must not share an axis name with a TensorCoreMesh when combined. The check enforces that sparse-core and tensor-core axis names are distinct so mapping does not become ambiguous.

Source

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

  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 [
        tpu_core.MemorySpace.VMEM_SHARED,
        tpu_core.MemorySpace.SMEM,
        tpu_core.MemorySpace.SEMAPHORE,
    ]

  @contextlib.contextmanager
  def tracing_context(self):
    yield

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Rename the ScalarSubcoreMesh axis (e.g. 'sc') so it differs from the TensorCoreMesh axis
2. Use distinct, prefixed axis names per core type by convention

Example fix

# before
scalar_mesh = pl_mosaic.ScalarSubcoreMesh(axis_name='cores', ...)
scalar_mesh.check_is_compatible_with(tensorcore_mesh)  # tensorcore_mesh.axis_name == 'cores'
# after
scalar_mesh = pl_mosaic.ScalarSubcoreMesh(axis_name='sc', ...)

assert scalar_mesh.axis_name != tc_mesh.axis_name, 'axis names must differ'

null

null

When it happens

Trigger: check_is_compatible_with(other_mesh) where other_mesh is a tpu_core.TensorCoreMesh with the same axis_name as the ScalarSubcoreMesh.

Common situations: Reusing a generic axis name like 'cores' or 'x' for both the TensorCore mesh and SparseCore mesh in a multi-core kernel.

Related errors

• {self} should have a different core axis name from the Tenso
• {self} should have a different subcore axis name from the Te
• 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.

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