What causes "Unsupported memory space: {space}" in jax-ml/jax?

Use mosaic_gpu_core.MemorySpace enum values instead of strings to avoid typos

jax-ml/jax · error · NotImplementedError

Unsupported memory space: {space}

Error message

Unsupported memory space: {space}

What it means

The interpreter can only allocate buffers in the GMEM, SMEM, TMEM, and REGS memory spaces. A Block/allocation requesting any other memory space string or enum value raises NotImplementedError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/interpret/jaxpr_interpret.py:84

    )

  @functools.cached_property
  def num_devices(self) -> int:
    """Computes the number of (SPMD) devices."""
    return math.prod(self.axis_sizes.values())


def _raise_if_unsupported_memory_space(
    space: mosaic_gpu_core.MemorySpace | None,
):
  # TODO(nrink): Support more memory spaces.
  if space is not None and space not in [
      mosaic_gpu_core.MemorySpace.GMEM,
      mosaic_gpu_core.MemorySpace.SMEM,
      mosaic_gpu_core.MemorySpace.TMEM,
      mosaic_gpu_core.MemorySpace.REGS,
  ]:
    raise NotImplementedError(f"Unsupported memory space: {space}")


def _raise_if_unsupported_collective_axes(
    mesh: mosaic_gpu_core.Mesh | None,
    is_collective_by_thread_cluster_axis: tuple[bool, ...],
):
  if not mesh or not mesh.thread_name:
    if any(is_collective_by_thread_cluster_axis):
      raise ValueError(
          "Requesting collective allocations, but no explicit thread axis"
          " specified."
      )
  else:
    # Note that the leading entries in `is_collective_by_thread__cluster_axis`
    # correspond to the cluster axes, while the last entry corresponds to the
    # thread axis within a block.
    *is_collective_by_cluster_axis, is_thread_axis_collective = (
        is_collective_by_thread_cluster_axis

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use mosaic_gpu_core.MemorySpace enum values instead of strings to avoid typos
Map TPU VMEM usage to SMEM/TMEM for GPU kernels
Update jax so new memory spaces are recognized
Check the printed space value in the message against the supported set

Example fix

# before
blk = mgpu.Buffer(memory_space='vmem', ...)  # TPU space on GPU
# after
blk = mgpu.Buffer(memory_space=mgpu.MemorySpace.SMEM, ...)

Defensive patterns

Strategy: type-guard

Validate before calling

SUPPORTED = {mgpu.MemorySpace.GMEM, mgpu.MemorySpace.SMEM, mgpu.MemorySpace.TMEM, mgpu.MemorySpace.REGS}
assert blk.memory_space in SUPPORTED

Type guard

def supported_space(s) -> bool:
    from jax._src.pallas.mosaic_gpu import core
    return s in {core.MemorySpace.GMEM, core.MemorySpace.SMEM, core.MemorySpace.TMEM, core.MemorySpace.REGS}

Prevention

Always use the MemorySpace enum, never raw strings
Replace TPU VMEM with SMEM/TMEM in GPU kernels

When it happens

Trigger: Declaring a Block or scratch allocation with memory_space set to a value outside {GMEM, SMEM, TMEM, REGS} (e.g. 'vmem', unknown strings, or new spaces) and running in interpret mode.

Common situations: Porting TPU kernels using VMEM to GPU without changing memory_space; typos in memory space strings; newer spaces unsupported by the installed jax interpreter.

Related errors

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