jax-ml/jax · error · ValueError

{self} should have a different subcore axis name from the Te

Error message

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

What it means

A VectorSubcoreMesh's subcore axis name (e.g. 'lane') must not collide with the TensorCoreMesh axis name when combined; the check runs after the core-axis check and enforces unique names across both.

Source

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

    return False

  def check_is_compatible_with(self, other_mesh):
    if isinstance(other_mesh, VectorSubcoreMesh):
      raise ValueError("You can't use two different VectorSubcoreMeshes.")
    elif isinstance(other_mesh, ScalarSubcoreMesh):
      if (other_mesh.axis_name == self.core_axis_name
          and other_mesh.num_cores == self.num_cores):
        return True
      raise ValueError(f"{self} should have the same core axis name and number"
                       f" of cores as the ScalarSubcoreMesh {other_mesh}.")
    elif isinstance(other_mesh, tpu_core.TensorCoreMesh):
      if self.core_axis_name == other_mesh.axis_name:
        raise ValueError(
            f"{self} should have a different core axis name from the"
            f" TensorCoreMesh {other_mesh}."
        )
      if self.subcore_axis_name == other_mesh.axis_name:
        raise ValueError(
            f"{self} should have a different subcore axis name from the"
            f" TensorCoreMesh {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,
        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

Rename the subcore_axis_name (e.g. 'sc_lane') so it is unique
Namespace all axis names by core type

Example fix

# before
vec = pl_mosaic.VectorSubcoreMesh(core_axis_name='sc', subcore_axis_name='lanes', ...)
# after
vec = pl_mosaic.VectorSubcoreMesh(core_axis_name='sc', subcore_axis_name='sc_lane', ...)

Defensive patterns

Strategy: validation

Validate before calling

assert vec.subcore_axis_name != tc_mesh.axis_name

Type guard

null

Try / catch

null

Prevention

Name subcore axes like 'sc_lane' to avoid collisions

When it happens

Trigger: check_is_compatible_with(tensor_core_mesh) where vector_mesh.subcore_axis_name == tensor_mesh.axis_name, e.g. both named 'lanes'.

Common situations: Choosing a subcore axis name that matches an existing tensor-core grid axis like 'x' or 'lanes'.

Related errors

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