jax-ml/jax · error · NotImplementedError

address_space not supported: {address_space}

Error message

address_space not supported: {address_space}

What it means

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.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:83

      f.write(content)
      f.write("\n")
  except OSError as e:
    logger.error("Failed to write output to %s: %s", filepath, e)
    # TODO(bchetioui): revisit whether this default of writing to stdout is the
    # right one. If we change it, we will have to change the corresponding C++
    # implementation as well.
    logger.error("Output will be written to stdout instead.")
    print(content)


def gpu_address_space_to_nvptx(address_space: gpu.AddressSpace) -> int:
  match address_space:
    case gpu.AddressSpace.Global:
      return 1
    case gpu.AddressSpace.Workgroup:
      return 3
    case _:
      raise NotImplementedError(f"address_space not supported: {address_space}")


WORKGROUP_NVPTX_ADDRESS_SPACE = gpu_address_space_to_nvptx(
    gpu.AddressSpace.Workgroup
)


def ptr_as_memref(ptr, memref_ty: ir.MemRefType):
  ptr_ty = llvm.PointerType(ptr.type)
  if ptr_ty.address_space != (get_memref_llvm_address_space(memref_ty) or 0):
    raise ValueError(
        f"Pointer address space {ptr_ty.address_space} does not match "
        f"memref memory space {memref_ty.memory_space}."
    )

  strides, offset = memref_ty.get_strides_and_offset()
  if offset != 0:
    raise ValueError("Non-zero offset is not supported for ptr_as_memref")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use gpu.AddressSpace.Global or gpu.AddressSpace.Workgroup for the memref's memory space
2. Remove explicit memory_space annotations so default lowering picks a supported space
3. Update to a newer Mosaic/JAX if the address space you need gained support

Example fix

# before
memref_ty = ir.MemRefType.get(shape, f32, memory_space=ir.Attribute.parse('#gpu.private'))
# after
from jax._src.interpreters.mlir import gpu
memref_ty = ir.MemRefType.get(shape, f32, memory_space=gpu.AddressSpace.Workgroup)

from jax._src.interpreters.mlir import gpu
assert address_space in (gpu.AddressSpace.Global, gpu.AddressSpace.Workgroup)

def is_supported_address_space(addr):
    return addr in (gpu.AddressSpace.Global, gpu.AddressSpace.Workgroup)

When it happens

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

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

Related errors

• Rule cannot handle an op with vector operands or results: {o
• Pointer address space {ptr_ty.address_space} does not match
• Too many dynamic shapes in the input. Mosaic currently only
• program id was requested but no grid was provided.
• Invalid axis {axis} for num_programs

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