jax-ml/jax · error · ValueError

You can't use two different TensorCoreMeshes.

Error message

You can't use two different TensorCoreMeshes.

What it means

Mosaic TPU kernels that use TensorCoreMesh (multi-core TPU execution) require a single consistent mesh. check_is_compatible_with rejects any combination of two different TensorCoreMesh instances, because a computation cannot span two independently-created meshes.

Source

Thrown at jax/_src/pallas/mosaic/core.py:426

  @property
  def default_memory_space(self) -> pallas_core.MemorySpace:
    return pallas_core.MemorySpace.ANY

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

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

  def discharges_effect(self, effect: jax_core.Effect) -> Literal[False]:
    del effect
    return False

  def check_is_compatible_with(self, other_mesh):
    if isinstance(other_mesh, TensorCoreMesh) and self != other_mesh:
      raise ValueError("You can't use two different TensorCoreMeshes.")
    # TODO: Add support for mpmd with SparseCore meshes.
    return super().check_is_compatible_with(other_mesh)

  @property
  def supported_memory_spaces(self) -> Sequence[Any]:
    return [
        MemorySpace.VMEM,
        MemorySpace.SMEM,
        MemorySpace.CMEM,
        MemorySpace.SEMAPHORE,
    ]

  @contextlib.contextmanager
  def tracing_context(self):
    yield


def create_tensorcore_mesh(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Create the TensorCoreMesh once (module-level or cached) and reuse it everywhere
2. Ensure equality: TensorCoreMesh compares by axis_name/num_cores — make parameters identical or, better, share the object
3. Pass the mesh explicitly through your code instead of constructing it ad hoc at call sites

Example fix

// before
def run(x):
  mesh = create_tensorcore_mesh('tc', num_cores=4)  # new mesh each call
  ...
// after
_MESH = create_tensorcore_mesh('tc', num_cores=4)
def run(x):
  mesh = _MESH  # single shared mesh

_MESH = None
def get_mesh():
    global _MESH
    if _MESH is None:
        _MESH = create_tensorcore_mesh('tc', num_cores=4)
    return _MESH

When it happens

Trigger: Calling a kernel created with create_tensorcore_mesh('mesh_a', num_cores=...) with arguments or effects produced under another TensorCoreMesh('mesh_b'), or nesting scoped computations under two different meshes.

Common situations: Creating the mesh once per function call (so each call gets a new, unequal mesh object) and reusing buffers/kernels across calls; mixing a global mesh with a locally-constructed one in a library refactor.

Related errors

• Memory space {self.memory_space} is not supported by mesh {s
• cannot specify both devices and num_cores
• Only gathers along the two minormost dimensions supported on
• When passing host local inputs to pjit, devices connected to
• Failed to find assignment for logical_axis_index {logical_ax

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