jax-ml/jax · error · NotImplementedError

MPMD map with more than one mesh requires scratch_type to ha

Error message

MPMD map with more than one mesh requires scratch_type to have a `core_type` specified, but {scratch_type=} is missing it.

What it means

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.

Source

Thrown at jax/_src/pallas/mpmd.py:919

    flat_scratch_types, scratch_tree = tree_util.tree_flatten(scratch_types)
    if len(meshes_and_fns) > 1:
      # TODO(rdyro): For MPMD with more than one mesh, come up with a better
      # solution for how to enforce core_type presence in scratch_shape.
      # TODO(rdyro): Check if we need to have a similar check for in-kernel
      # allocations (e.g., run_scoped, empty_ref) or can we assume the
      # core_type is inherited from the caller (we then need the core_type in
      # the caller context during tracing).
      # TODO(rdyro): Also check inputs and outputs for core type.
      for scratch_type in flat_scratch_types:
        from jax._src.pallas.mosaic import core as tpu_core  # pyrefly: ignore[missing-import]

        if not isinstance(
            scratch_type.memory_space, pallas_core.CoreMemorySpace
        ) and scratch_type.memory_space not in (
            tpu_core.MemorySpace.HBM,
            tpu_core.MemorySpace.VMEM_SHARED,
        ):
          raise NotImplementedError(
              "MPMD map with more than one mesh requires scratch_type to have"
              f" a `core_type` specified, but {scratch_type=} is missing it."
          )

    # Kernels may have Refs that belong to external meshes (usually for
    # async kernels). For example, the SC ScalarSubcore may have a Reference
    # to a TC semaphore that it is signaling. There is no explicit TC mesh as
    # part of the user-provided meshes, and are instead snuck in via the aval.
    for aval in [*flat_avals, *flat_out_avals, *flat_scratch_types]:
      if (
          isinstance(aval, jax_core.ShapedArray)
          and isinstance(aval.memory_space, pallas_core.CoreMemorySpace)
          and aval.memory_space.mesh not in it.chain(meshes, external_meshes)
      ):
        external_meshes.append(aval.memory_space.mesh)

    all_meshes = (*meshes, *external_meshes)
    # Check that meshes are compatible with each other (e.g, have a consistent

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. 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
2. Use tpu_core.MemorySpace.VMEM_SHARED or HBM which are valid without an explicit core binding
3. Reduce the setup to a single mesh if per-core scratch is not required

Example fix

# before
scratch_type = tpu_core.BlockType(dtype=jnp.float32, block_shape=(128,))
mpmd_map(kernel, ..., scratch_type=scratch_type)  # multiple meshes
# after
scratch_type = tpu_core.BlockType(
    dtype=jnp.float32, block_shape=(128,),
    memory_space=pallas_core.CoreMemorySpace(
        tpu_core.MemorySpace.VMEM, core_type=my_core_type))

from jax._src import pallas_core
from jax.experimental import tpu_core

def scratch_ok_for_multimesh(scratch_type, num_meshes):
    if num_meshes <= 1:
        return True
    ms = scratch_type.memory_space
    return isinstance(ms, pallas_core.CoreMemorySpace) or ms in (
        tpu_core.MemorySpace.HBM, tpu_core.MemorySpace.VMEM_SHARED)

When it happens

Trigger: 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).

Common situations: 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.

Related errors

• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s
• Acc ref must be at least 2D, got shape {shape}
• Acc ref dtype must be float32 or int32, got {dtype}
• Accumulators are not available on TPU {info.chip_version}

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