{"record":{"id":"b68f1e44f504b565","repo":"jax-ml/jax","slug":"could-not-find-memory-addressable-by-device-s-de","errorCode":null,"errorMessage":"Could not find memory addressable by device %s. Device %s can address the following memory kinds: %s. Got memory kind: %s","messagePattern":"Could not find memory addressable by device (.+?)\\. Device (.+?) can address the following memory kinds: (.+?)\\. Got memory kind: (.+?)","errorType":"validation","errorClass":"nb::value_error","httpStatus":null,"severity":"error","filePath":"jaxlib/sharding.cc","lineNumber":106,"sourceCode":"    for (nb::handle supported_memory_kind : *supported_memory_kinds) {\n      if (supported_memory_kind.equal(memory_kind)) {\n        return memory_kind;\n      }\n    }\n    auto addressable_device_list =\n        PyDeviceList::AddressableDeviceList(device_list);\n    if (addressable_device_list->Len() == 0) {\n      // If the device list is not addressable, we can't check if the memory\n      // kind is supported, so we assume it is.\n      return memory_kind;\n    }\n    nb::object device_kind =\n        addressable_device_list->GetItem(0).attr(\"device_kind\");\n    std::string_view device_kind_str = nb::cast<std::string_view>(device_kind);\n    auto py_str_formatter = [](std::string* out, nb::handle h) {\n      *out += nb::cast<std::string_view>(nb::str(h));\n    };\n    throw nb::value_error(\n        absl::StrCat(\n            \"Could not find memory addressable by device \", device_kind_str,\n            \". Device \", device_kind_str,\n            \" can address the following memory kinds: \",\n            absl::StrJoin(*supported_memory_kinds, \", \", py_str_formatter),\n            \". Got memory kind: \", nb::cast<std::string_view>(memory_kind))\n            .c_str());\n  }\n  // If memory kind is None, canonicalize to default memory.\n  absl::StatusOr<nb::object> default_memory_kind =\n      PyDeviceList::DefaultMemoryKind(device_list);\n  if (!default_memory_kind.ok()) {\n    return nb::none();\n  }\n  return *std::move(default_memory_kind);\n}\n\n// This list is to check for valid memory kinds when an AbstractMesh is passed","sourceCodeStart":88,"sourceCodeEnd":124,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jaxlib/sharding.cc#L88-L124","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Read the error message: it lists valid kinds for your device; use one of those","Query supported kinds at runtime from the device (jax.devices()[0].addressable_memories or backend attributes) instead of hardcoding","Make memory_kind configurable or default to None so JAX picks the device default","Fix typos: common kinds are 'device' (HBM), 'pinned_host_memory', 'unpinned_host_memory'"],"exampleFix":"# before\nsh = jax.sharding.NamedSharding(jax.devices()[0], memory_kind='hbm')\n\n# after\nsh = jax.sharding.NamedSharding(jax.devices()[0], memory_kind='device')  # or omit memory_kind","handlingStrategy":"type-guard","validationCode":"dev = jax.devices()[0]\nkinds = {m.kind for m in getattr(dev, 'addressable_memories', [])}\nif memory_kind is not None and memory_kind not in kinds:\n    raise ValueError(f'{memory_kind} not in {kinds}; pick one of these')\nsh = jax.sharding.NamedSharding(dev, memory_kind=memory_kind)","typeGuard":"def memory_kind_supported(device, kind: str | None) -> bool:\n    if kind is None:\n        return True\n    return kind in {m.kind for m in device.addressable_memories}","tryCatchPattern":"try:\n    sh = jax.sharding.NamedSharding(dev, memory_kind=kind)\nexcept ValueError as e:\n    if 'Could not find memory' in str(e):\n        sh = jax.sharding.NamedSharding(dev)  # default memory kind\n    else:\n        raise","preventionTips":["Never hardcode 'hbm'; the device kind is named 'device'","Derive memory kinds from device.addressable_memories at runtime","Default memory_kind=None in reusable code"],"tags":["jax","sharding","memory-kind","tpu"],"backgroundTag":"invalid-device-memory-kind","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}