{"record":{"id":"60922244477ce671","repo":"jax-ml/jax","slug":"mpmd-map-with-more-than-one-mesh-requires-scratch","errorCode":null,"errorMessage":"MPMD map with more than one mesh requires scratch_type to have a `core_type` specified, but {scratch_type=} is missing it.","messagePattern":"MPMD map with more than one mesh requires scratch_type to have a `core_type` specified, but (.+?) is missing it\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mpmd.py","lineNumber":919,"sourceCode":"    flat_scratch_types, scratch_tree = tree_util.tree_flatten(scratch_types)\n    if len(meshes_and_fns) > 1:\n      # TODO(rdyro): For MPMD with more than one mesh, come up with a better\n      # solution for how to enforce core_type presence in scratch_shape.\n      # TODO(rdyro): Check if we need to have a similar check for in-kernel\n      # allocations (e.g., run_scoped, empty_ref) or can we assume the\n      # core_type is inherited from the caller (we then need the core_type in\n      # the caller context during tracing).\n      # TODO(rdyro): Also check inputs and outputs for core type.\n      for scratch_type in flat_scratch_types:\n        from jax._src.pallas.mosaic import core as tpu_core  # pyrefly: ignore[missing-import]\n\n        if not isinstance(\n            scratch_type.memory_space, pallas_core.CoreMemorySpace\n        ) and scratch_type.memory_space not in (\n            tpu_core.MemorySpace.HBM,\n            tpu_core.MemorySpace.VMEM_SHARED,\n        ):\n          raise NotImplementedError(\n              \"MPMD map with more than one mesh requires scratch_type to have\"\n              f\" a `core_type` specified, but {scratch_type=} is missing it.\"\n          )\n\n    # Kernels may have Refs that belong to external meshes (usually for\n    # async kernels). For example, the SC ScalarSubcore may have a Reference\n    # to a TC semaphore that it is signaling. There is no explicit TC mesh as\n    # part of the user-provided meshes, and are instead snuck in via the aval.\n    for aval in [*flat_avals, *flat_out_avals, *flat_scratch_types]:\n      if (\n          isinstance(aval, jax_core.ShapedArray)\n          and isinstance(aval.memory_space, pallas_core.CoreMemorySpace)\n          and aval.memory_space.mesh not in it.chain(meshes, external_meshes)\n      ):\n        external_meshes.append(aval.memory_space.mesh)\n\n    all_meshes = (*meshes, *external_meshes)\n    # Check that meshes are compatible with each other (e.g, have a consistent","sourceCodeStart":901,"sourceCodeEnd":937,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mpmd.py#L901-L937","documentation":"When a Pallas MPMD map spans more than one mesh, scratch memory must be allocated per core, so the scratch_type must be tied to a specific core via a CoreMemorySpace (or the shared VMEM space). If scratch_type.memory_space is a generic memory space (not core-specific) such as plain HBM/VMEM without a core_type, MPMD cannot tell which device's core owns the scratch and raises NotImplementedError.","triggerScenarios":"Calling mpmd_map with multiple meshes while passing a scratch_type whose memory_space is not an instance of pallas_core.CoreMemorySpace and not one of (tpu_core.MemorySpace.HBM, tpu_core.MemorySpace.VMEM_SHARED).","commonSituations":"Migrating single-mesh Pallas code (where scratch_type with a plain memory space was fine) to multi-mesh/multi-host MPMD; using tpu_core.BlockType or scratch specs written for the single-core API.","solutions":["Give the scratch type a core-specific memory space, e.g. wrap it with pallas_core.CoreMemorySpace(tpu_core.MemorySpace.VMEM, core_type=...) or set the core_type on the scratch type","Use tpu_core.MemorySpace.VMEM_SHARED or HBM which are valid without an explicit core binding","Reduce the setup to a single mesh if per-core scratch is not required"],"exampleFix":"# before\nscratch_type = tpu_core.BlockType(dtype=jnp.float32, block_shape=(128,))\nmpmd_map(kernel, ..., scratch_type=scratch_type)  # multiple meshes\n# after\nscratch_type = tpu_core.BlockType(\n    dtype=jnp.float32, block_shape=(128,),\n    memory_space=pallas_core.CoreMemorySpace(\n        tpu_core.MemorySpace.VMEM, core_type=my_core_type))","handlingStrategy":"validation","validationCode":"from jax._src import pallas_core\nfrom jax.experimental import tpu_core\n\ndef scratch_ok_for_multimesh(scratch_type, num_meshes):\n    if num_meshes <= 1:\n        return True\n    ms = scratch_type.memory_space\n    return isinstance(ms, pallas_core.CoreMemorySpace) or ms in (\n        tpu_core.MemorySpace.HBM, tpu_core.MemorySpace.VMEM_SHARED)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["When using multiple meshes, always specify a core_type (CoreMemorySpace) or VMEM_SHARED/HBM for scratch","Encapsulate scratch_type construction in one helper so MPMD-safe defaults are always used"],"tags":["jax","pallas","mpmd","tpu","scratch-memory","multi-mesh"],"backgroundTag":"missing-core-type-in-scratch-config","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}