{"record":{"id":"166855357bd58f37","repo":"jax-ml/jax","slug":"unsupported-memory-space-space","errorCode":null,"errorMessage":"Unsupported memory space: {space}","messagePattern":"Unsupported memory space: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/interpret/jaxpr_interpret.py","lineNumber":84,"sourceCode":"    )\n\n  @functools.cached_property\n  def num_devices(self) -> int:\n    \"\"\"Computes the number of (SPMD) devices.\"\"\"\n    return math.prod(self.axis_sizes.values())\n\n\ndef _raise_if_unsupported_memory_space(\n    space: mosaic_gpu_core.MemorySpace | None,\n):\n  # TODO(nrink): Support more memory spaces.\n  if space is not None and space not in [\n      mosaic_gpu_core.MemorySpace.GMEM,\n      mosaic_gpu_core.MemorySpace.SMEM,\n      mosaic_gpu_core.MemorySpace.TMEM,\n      mosaic_gpu_core.MemorySpace.REGS,\n  ]:\n    raise NotImplementedError(f\"Unsupported memory space: {space}\")\n\n\ndef _raise_if_unsupported_collective_axes(\n    mesh: mosaic_gpu_core.Mesh | None,\n    is_collective_by_thread_cluster_axis: tuple[bool, ...],\n):\n  if not mesh or not mesh.thread_name:\n    if any(is_collective_by_thread_cluster_axis):\n      raise ValueError(\n          \"Requesting collective allocations, but no explicit thread axis\"\n          \" specified.\"\n      )\n  else:\n    # Note that the leading entries in `is_collective_by_thread__cluster_axis`\n    # correspond to the cluster axes, while the last entry corresponds to the\n    # thread axis within a block.\n    *is_collective_by_cluster_axis, is_thread_axis_collective = (\n        is_collective_by_thread_cluster_axis","sourceCodeStart":66,"sourceCodeEnd":102,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/interpret/jaxpr_interpret.py#L66-L102","documentation":"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.","triggerScenarios":"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.","commonSituations":"Porting TPU kernels using VMEM to GPU without changing memory_space; typos in memory space strings; newer spaces unsupported by the installed jax interpreter.","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"],"exampleFix":"# before\nblk = mgpu.Buffer(memory_space='vmem', ...)  # TPU space on GPU\n# after\nblk = mgpu.Buffer(memory_space=mgpu.MemorySpace.SMEM, ...)","handlingStrategy":"type-guard","validationCode":"SUPPORTED = {mgpu.MemorySpace.GMEM, mgpu.MemorySpace.SMEM, mgpu.MemorySpace.TMEM, mgpu.MemorySpace.REGS}\nassert blk.memory_space in SUPPORTED","typeGuard":"def supported_space(s) -> bool:\n    from jax._src.pallas.mosaic_gpu import core\n    return s in {core.MemorySpace.GMEM, core.MemorySpace.SMEM, core.MemorySpace.TMEM, core.MemorySpace.REGS}","tryCatchPattern":null,"preventionTips":["Always use the MemorySpace enum, never raw strings","Replace TPU VMEM with SMEM/TMEM in GPU kernels"],"tags":["pallas","mosaic-gpu","memory-space","interpret-mode","not-implemented"],"backgroundTag":"unsupported-operation-in-interpreter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}