[JAX] Add a test for multiprocess shard_map in McJAX with non-participating hosts. Update handling of device memories in PyDeviceList to support this.

PiperOrigin-RevId: 753749996

Author: emilyfertig

jax/_src/array.py

@@ -636,7 +636,7 @@ def _value(self) -> np.ndarray:
     self._check_if_deleted()
 
     if self._npy_value is None:
-      if self.is_fully_replicated:
+      if self.is_fully_replicated and self.sharding.addressable_devices:
         npy_value, did_copy = self._single_device_array_to_np_array_did_copy()
         npy_value.flags.writeable = False
         if did_copy:

jaxlib/py_array.cc

@@ -1528,6 +1528,9 @@ int PyArray_bf_getbuffer(PyObject* exporter, Py_buffer* view, int flags) {
     absl::Span<const std::shared_ptr<PjRtBuffer>> buffers =
         array->pjrt_buffers();
 
+    if (buffers.empty()) {
+      return InvalidArgument("Array has no buffers.");
+    }
     PjRtBuffer& buffer = *buffers.front();
     if (!buffer.IsOnCpu()) {
       return InvalidArgument(

jaxlib/py_device_list.cc

@@ -366,31 +366,19 @@ void PyDeviceList::PopulateMemoryKindInfo() {
     throw nb::value_error("Unrecognized DeviceList type");
   }
   MemoryKindInfo info;
-  xla::ifrt::Device* addressable_device = nullptr;
-  const int process_index = py_client_ ? py_client_->process_index() : 0;
-  for (xla::ifrt::Device* device : std::get<0>(device_list_)->devices()) {
-    if (device->ProcessIndex() == process_index) {
-      addressable_device = device;
-      break;
-    }
-  }
-  if (addressable_device == nullptr) {
-    info.default_memory_kind = nb::none();
-    memory_kind_info_ = std::move(info);
-    return;
-  }
+  xla::ifrt::Device* device = std::get<0>(device_list_)->devices()[0];
 
-  auto default_memory = addressable_device->DefaultMemory();
+  auto default_memory = device->DefaultMemory();
   if (!default_memory.ok()) {
     // Cache the error.
     memory_kind_info_ = default_memory.status();
     return;
   }
   info.default_memory_kind = nb::cast(*(*default_memory)->Kind().memory_kind());
   nb::tuple memory_kinds =
-      nb::steal<nb::tuple>(PyTuple_New(addressable_device->Memories().size()));
-  for (size_t i = 0; i < addressable_device->Memories().size(); ++i) {
-    auto* memory = addressable_device->Memories()[i];
+      nb::steal<nb::tuple>(PyTuple_New(device->Memories().size()));
+  for (size_t i = 0; i < device->Memories().size(); ++i) {
+    auto* memory = device->Memories()[i];
     nb::str s = nb::str(memory->Kind().memory_kind()->data(),
                         memory->Kind().memory_kind()->size());
     PyTuple_SET_ITEM(memory_kinds.ptr(), i, s.release().ptr());