Internal Change. DO NOT SUBMIT: Remove this line after the baseline is

submitted.

PiperOrigin-RevId: 718995567

jax/_src/xla_bridge.py

@@ -1015,7 +1015,9 @@ def _init_backend(platform: str) -> xla_client.Client:
   # factories instead of returning None.
   if backend is None:
     raise RuntimeError(f"Could not initialize backend '{platform}'")
-  if backend.device_count() == 0:
+  # TODO(b/356678989): Only check `backend.device_count()` when it counts
+  # CPU-only devices.
+  if backend.device_count() == 0 and len(backend._get_all_devices()) == 0:
     raise RuntimeError(f"Backend '{platform}' provides no devices.")
   util.distributed_debug_log(("Initialized backend", backend.platform),
                              ("process_index", backend.process_index()),

tests/xla_bridge_test.py

@@ -217,9 +217,15 @@ def device_count(self):
     def process_index(self):
       return 0
 
+    def devices(self):
+      return []
+
     def local_devices(self):
       return []
 
+    def _get_all_devices(self):
+      return self.devices()
+
   def _register_factory(self, platform: str, priority, device_count=1,
                         assert_used_at_most_once=False, experimental=False):
     if assert_used_at_most_once: