jax-ml/jax · error · nb::value_error

Could not find memory addressable by device %s. Device %s ca

Error message

Could not find memory addressable by device %s. Device %s can address the following memory kinds: %s. Got memory kind: %s

What it means

Raised by CheckAndCanonicalizeMemoryKind when constructing a NamedSharding/SingleDeviceSharding/GSPMDSharding with an explicit memory_kind (e.g. 'hbm', 'unpinned_host_memory') that none of the sharding's addressable devices can access. The message lists the memory kinds the device actually supports.

Source

Thrown at jaxlib/sharding.cc:106

    for (nb::handle supported_memory_kind : *supported_memory_kinds) {
      if (supported_memory_kind.equal(memory_kind)) {
        return memory_kind;
      }
    }
    auto addressable_device_list =
        PyDeviceList::AddressableDeviceList(device_list);
    if (addressable_device_list->Len() == 0) {
      // If the device list is not addressable, we can't check if the memory
      // kind is supported, so we assume it is.
      return memory_kind;
    }
    nb::object device_kind =
        addressable_device_list->GetItem(0).attr("device_kind");
    std::string_view device_kind_str = nb::cast<std::string_view>(device_kind);
    auto py_str_formatter = [](std::string* out, nb::handle h) {
      *out += nb::cast<std::string_view>(nb::str(h));
    };
    throw nb::value_error(
        absl::StrCat(
            "Could not find memory addressable by device ", device_kind_str,
            ". Device ", device_kind_str,
            " can address the following memory kinds: ",
            absl::StrJoin(*supported_memory_kinds, ", ", py_str_formatter),
            ". Got memory kind: ", nb::cast<std::string_view>(memory_kind))
            .c_str());
  }
  // If memory kind is None, canonicalize to default memory.
  absl::StatusOr<nb::object> default_memory_kind =
      PyDeviceList::DefaultMemoryKind(device_list);
  if (!default_memory_kind.ok()) {
    return nb::none();
  }
  return *std::move(default_memory_kind);
}

// This list is to check for valid memory kinds when an AbstractMesh is passed

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Read the error message: it lists valid kinds for your device; use one of those
2. Query supported kinds at runtime from the device (jax.devices()[0].addressable_memories or backend attributes) instead of hardcoding
3. Make memory_kind configurable or default to None so JAX picks the device default
4. Fix typos: common kinds are 'device' (HBM), 'pinned_host_memory', 'unpinned_host_memory'

Example fix

# before
sh = jax.sharding.NamedSharding(jax.devices()[0], memory_kind='hbm')

# after
sh = jax.sharding.NamedSharding(jax.devices()[0], memory_kind='device')  # or omit memory_kind

dev = jax.devices()[0]
kinds = {m.kind for m in getattr(dev, 'addressable_memories', [])}
if memory_kind is not None and memory_kind not in kinds:
    raise ValueError(f'{memory_kind} not in {kinds}; pick one of these')
sh = jax.sharding.NamedSharding(dev, memory_kind=memory_kind)

def memory_kind_supported(device, kind: str | None) -> bool:
    if kind is None:
        return True
    return kind in {m.kind for m in device.addressable_memories}

try:
    sh = jax.sharding.NamedSharding(dev, memory_kind=kind)
except ValueError as e:
    if 'Could not find memory' in str(e):
        sh = jax.sharding.NamedSharding(dev)  # default memory kind
    else:
        raise

When it happens

Trigger: Passing memory_kind to a sharding where the target device doesn't expose that kind, e.g. memory_kind='hbm' on a device that only supports 'vmem', or a typo'd kind, or an OpenXLA backend with restricted addressable memory.

Common situations: Writing device-agnostic JAX code that hardcodes memory_kind='hbm' then running on a TPU/CPU/plugin backend; typos in memory kind strings; backend plugins that don't register the expected memory kinds; using unpinned_host_memory with devices lacking host memory reporting.

Related errors

• Partitioned callback not implemented on {platform} backend.
• sharding with memory_kind is not allowed. Please use `jax.de
• Custom Partitioning rules must return Sharding.
• Subclasses should implement this method
• Mosaic kernels cannot be automatically partitioned. Please w

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