{"record":{"id":"6ecf8149a8c90ca0","repo":"jax-ml/jax","slug":"address-space-not-supported-address-space","errorCode":null,"errorMessage":"address_space not supported: {address_space}","messagePattern":"address_space not supported: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":83,"sourceCode":"      f.write(content)\n      f.write(\"\\n\")\n  except OSError as e:\n    logger.error(\"Failed to write output to %s: %s\", filepath, e)\n    # TODO(bchetioui): revisit whether this default of writing to stdout is the\n    # right one. If we change it, we will have to change the corresponding C++\n    # implementation as well.\n    logger.error(\"Output will be written to stdout instead.\")\n    print(content)\n\n\ndef gpu_address_space_to_nvptx(address_space: gpu.AddressSpace) -> int:\n  match address_space:\n    case gpu.AddressSpace.Global:\n      return 1\n    case gpu.AddressSpace.Workgroup:\n      return 3\n    case _:\n      raise NotImplementedError(f\"address_space not supported: {address_space}\")\n\n\nWORKGROUP_NVPTX_ADDRESS_SPACE = gpu_address_space_to_nvptx(\n    gpu.AddressSpace.Workgroup\n)\n\n\ndef ptr_as_memref(ptr, memref_ty: ir.MemRefType):\n  ptr_ty = llvm.PointerType(ptr.type)\n  if ptr_ty.address_space != (get_memref_llvm_address_space(memref_ty) or 0):\n    raise ValueError(\n        f\"Pointer address space {ptr_ty.address_space} does not match \"\n        f\"memref memory space {memref_ty.memory_space}.\"\n    )\n\n  strides, offset = memref_ty.get_strides_and_offset()\n  if offset != 0:\n    raise ValueError(\"Non-zero offset is not supported for ptr_as_memref\")","sourceCodeStart":65,"sourceCodeEnd":101,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L65-L101","documentation":"Raised by gpu_address_space_to_nvptx when mapping an MLIR gpu.AddressSpace to its NVPTX LLVM address-space number. Only Global (->1) and Workgroup (->3) are implemented; any other address space (e.g. Private or a custom integer attribute) has no NVPTX mapping in this helper.","triggerScenarios":"Passing gpu.AddressSpace.Private (or an address space other than Global/Workgroup) to gpu_address_space_to_nvptx, or to helpers built on it like get_memref_llvm_address_space or workgroup_ptr_ty — e.g. lowering a memref whose memory_space is set to #gpu.private or a raw integer space.","commonSituations":"Hand-annotating memref memory spaces in custom lowerings; Mosaic version changes that added/renamed address-space enums; constructing IR from other dialects that emit non-standard memory spaces.","solutions":["Use gpu.AddressSpace.Global or gpu.AddressSpace.Workgroup for the memref's memory space","Remove explicit memory_space annotations so default lowering picks a supported space","Update to a newer Mosaic/JAX if the address space you need gained support"],"exampleFix":"# before\nmemref_ty = ir.MemRefType.get(shape, f32, memory_space=ir.Attribute.parse('#gpu.private'))\n# after\nfrom jax._src.interpreters.mlir import gpu\nmemref_ty = ir.MemRefType.get(shape, f32, memory_space=gpu.AddressSpace.Workgroup)","handlingStrategy":"type-guard","validationCode":"from jax._src.interpreters.mlir import gpu\nassert address_space in (gpu.AddressSpace.Global, gpu.AddressSpace.Workgroup)","typeGuard":"def is_supported_address_space(addr):\n    return addr in (gpu.AddressSpace.Global, gpu.AddressSpace.Workgroup)","tryCatchPattern":null,"preventionTips":["Restrict memref memory_space annotations to Global/Workgroup in custom lowerings","Don't propagate foreign dialect address spaces into Mosaic pointer utilities"],"tags":["gpu","mosaic","address-space","nvptx","lowering"],"backgroundTag":"unsupported-address-space","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}