{"record":{"id":"13ca4c703c451b3a","repo":"jax-ml/jax","slug":"you-can-t-use-two-different-tensorcoremeshes","errorCode":null,"errorMessage":"You can't use two different TensorCoreMeshes.","messagePattern":"You can't use two different TensorCoreMeshes\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/core.py","lineNumber":426,"sourceCode":"  @property\n  def default_memory_space(self) -> pallas_core.MemorySpace:\n    return pallas_core.MemorySpace.ANY\n\n  @property\n  def shape(self):\n    return collections.OrderedDict({self.axis_name: self.num_cores})\n\n  @property\n  def dimension_semantics(self) -> Sequence[DimensionSemantics]:\n    return [GridDimensionSemantics.PARALLEL]\n\n  def discharges_effect(self, effect: jax_core.Effect) -> Literal[False]:\n    del effect\n    return False\n\n  def check_is_compatible_with(self, other_mesh):\n    if isinstance(other_mesh, TensorCoreMesh) and self != other_mesh:\n      raise ValueError(\"You can't use two different TensorCoreMeshes.\")\n    # TODO: Add support for mpmd with SparseCore meshes.\n    return super().check_is_compatible_with(other_mesh)\n\n  @property\n  def supported_memory_spaces(self) -> Sequence[Any]:\n    return [\n        MemorySpace.VMEM,\n        MemorySpace.SMEM,\n        MemorySpace.CMEM,\n        MemorySpace.SEMAPHORE,\n    ]\n\n  @contextlib.contextmanager\n  def tracing_context(self):\n    yield\n\n\ndef create_tensorcore_mesh(","sourceCodeStart":408,"sourceCodeEnd":444,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/core.py#L408-L444","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Create the TensorCoreMesh once (module-level or cached) and reuse it everywhere","Ensure equality: TensorCoreMesh compares by axis_name/num_cores — make parameters identical or, better, share the object","Pass the mesh explicitly through your code instead of constructing it ad hoc at call sites"],"exampleFix":"// before\ndef run(x):\n  mesh = create_tensorcore_mesh('tc', num_cores=4)  # new mesh each call\n  ...\n// after\n_MESH = create_tensorcore_mesh('tc', num_cores=4)\ndef run(x):\n  mesh = _MESH  # single shared mesh","handlingStrategy":"validation","validationCode":"_MESH = None\ndef get_mesh():\n    global _MESH\n    if _MESH is None:\n        _MESH = create_tensorcore_mesh('tc', num_cores=4)\n    return _MESH","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Create meshes once and cache at module level","Pass the mesh explicitly rather than reconstructing per call"],"tags":["jax","pallas","tpu","mesh","tensorcore","multi-core"],"backgroundTag":"inconsistent-mesh-configuration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}