jax-ml/jax · error · ValueError

Memory space {self.memory_space} is not supported by mesh {s

Error message

Memory space {self.memory_space} is not supported by mesh {self.mesh}

What it means

Pallas meshes declare which memory spaces (e.g. HBM, VMEM/SMEM per backend) they support. AbstractMemorySpec.__post_init__ (grid/mesh helpers) validates that the requested memory_space appears in mesh.supported_memory_spaces and rejects mismatches, preventing kernels from allocating in memory the target doesn't offer.

Source

Thrown at jax/_src/pallas/core.py:308

  def __call__(self, shape: tuple[int, ...], dtype: jnp.dtype):
    # A convenience function for constructing MemoryRef types of ShapedArrays.
    return self.from_type(jax_core.ShapedArray(shape, dtype))

  def __str__(self) -> str:
    return self.value


@dataclasses.dataclass(frozen=True)
class CoreMemorySpace:
  """A memory space tied to a Pallas mesh."""

  memory_space: Any
  mesh: Mesh

  def __post_init__(self):
    if not self.memory_space in self.mesh.supported_memory_spaces:
      raise ValueError(
          f"Memory space {self.memory_space} is not supported by mesh"
          f" {self.mesh}"
      )

  def __call__(self, shape: Sequence[int], dtype: jnp.dtype[Any]):
    return MemoryRef(jax_core.ShapedArray(tuple(shape), dtype), self)

  def __str__(self) -> str:
    return f"{self.memory_space}@{self.mesh.core_type}"

  def __repr__(self) -> str:
    return f"{self.memory_space!r}@{self.mesh.core_type!r}"

  @property
  def name(self) -> Any:
    return f"{self.memory_space}@{self.mesh.core_type.name}"

  @property

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Print mesh.supported_memory_spaces and use one of those exact values
2. Use backend-provided constants (e.g. tpu.VMEM / triton.SMEM) instead of raw strings
3. Ensure the mesh was constructed for the backend you target

Example fix

# before
spec = pl.core.AbstractMemorySpec(memory_space='vmem', mesh=mesh)  # GPU mesh
# after
from jax.experimental.pallas import triton as pl_gpu
spec = pl.core.AbstractMemorySpec(memory_space=pl_gpu.SMEM, mesh=mesh)

assert memory_space in mesh.supported_memory_spaces, (
    f'{memory_space} not in {mesh.supported_memory_spaces}')

When it happens

Trigger: Creating a Mesh-related memory spec with a memory space string/object not in the mesh's supported list — e.g. requesting 'vmem' on a mesh whose backend only supports HBM/SMEM names, or mixing TPU (VMEM) and GPU (SMEM) memory-space constants.

Common situations: Porting a TPU Pallas kernel (using VMEM) to GPU where the space is named differently; typos in memory space names; backend version differences in supported spaces.

Related errors

• Compiler params for platform {platform} cannot be used for {
• You can't use two different TensorCoreMeshes.
• cannot specify both devices and num_cores
• Unsupported core type: {core_type}
• Invalid memory space: {memory_space!r}

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