jax-ml/jax · error · ValueError

You can't use two different VectorSubcoreMeshes.

Error message

You can't use two different VectorSubcoreMeshes.

What it means

Mirror of the scalar case: VectorSubcoreMesh.check_is_compatible_with rejects pairing with another VectorSubcoreMesh. Only one vector subcore mesh may participate, since it fully describes the SparseCore subcore layout.

Source

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

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

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

  def discharges_effect(self, effect):
    del effect  # Unused.
    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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Share a single VectorSubcoreMesh instance
2. If one side should be scalar, use ScalarSubcoreMesh instead

Example fix

# before
vec_a.check_is_compatible_with(vec_b)  # both VectorSubcoreMesh
# after
vec_a.check_is_compatible_with(scalar_mesh)  # ScalarSubcoreMesh with matching axis/cores

def is_vector_mesh(m): return isinstance(m, pl_mosaic.VectorSubcoreMesh)

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

null

When it happens

Trigger: Calling VectorSubcoreMesh.check_is_compatible_with with another VectorSubcoreMesh instance.

Common situations: Two modules each creating their own VectorSubcoreMesh and later combining kernels/meshes.

Related errors

• You can't use two different ScalarSubcoreMeshes.
• {self} should have the same core axis name and number of cor
• {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/fbd11297189b99f2. Report an issue: GitHub.