Repository: google/jax

Issue:
jax.devices returns only the devices from the default backend
The documentation of jax.devices says that it returns all devices. This is not true in multi-backend setups, it returns only the devices from the default backend. This was confusing to me initially, and I also encountered this confusion in issue #2785 

The simplest change would be to the documentation: if no `backend` is specified then only the devices on the default backend are returned (along with an explanation of what is the default backend). 

A better change may be though to say that it returns all devices, in the order of backend priority, such that `devices()[0]` is the same as now. This may break some code though. 




Discussion:
I searched in Google3 and could not find any code that uses the backend parameter for `jax.devices` other than tests.

Resolving patch:
diff --git a/jax/lib/xla_bridge.py b/jax/lib/xla_bridge.py
--- a/jax/lib/xla_bridge.py
+++ b/jax/lib/xla_bridge.py
@@ -162,7 +162,7 @@ def get_device_backend(device=None):
   return get_backend(platform)
 
 
-def device_count(backend=None):
+def device_count(backend: str = None):
   """Returns the total number of devices.
 
   On most platforms, this is the same as ``local_device_count()``. However, on
@@ -179,19 +179,22 @@ def device_count(backend=None):
   return int(get_backend(backend).device_count())
 
 
-def local_device_count(backend=None):
+def local_device_count(backend: str =None):
   """Returns the number of devices on this host."""
   return int(get_backend(backend).local_device_count())
 
 
-def devices(backend=None):
-  """Returns a list of all devices.
+def devices(backend: str = None):
+  """Returns a list of all devices for a given backend.
 
-  Each device is represented by a subclass of Device (e.g. CpuDevice,
-  GpuDevice). The length of the returned list is equal to
-  ``device_count()``. Local devices can be identified by comparing
+  Each device is represented by a subclass of ``Device`` (e.g. ``CpuDevice``,
+  ``GpuDevice``). The length of the returned list is equal to
+  ``device_count(backend)``. Local devices can be identified by comparing
   ``Device.host_id`` to ``host_id()``.
 
+  If ``backend`` is ``None``, returns all the devices from the default backend.
+  The default backend is generally 'gpu' or 'tpu' if available, otherwise 'cpu'.
+
   Args:
     backend: This is an experimental feature and the API is likely to change.
       Optional, a string representing the xla backend. 'cpu', 'gpu', or 'tpu'.
@@ -202,14 +205,28 @@ def devices(backend=None):
   return get_backend(backend).devices()
 
 
-def local_devices(host_id=None, backend=None):
-  """Returns a list of devices local to a given host (this host by default)."""
+def local_devices(host_id: int = None, backend: str = None):
+  """Like ``devices``, but only returns devices local to a given host.
+
+  If ``host_id`` is ``None``, returns devices local to this host.
+
+  Args:
+    host_id: the integer ID of the host. Host IDs can be retrieved via
+      ``host_ids()``.
+    backend: This is an experimental feature and the API is likely to change.
+      Optional, a string representing the xla backend. 'cpu', 'gpu', or 'tpu'.
+
+  Returns:
+    List of Device subclasses.
+  """
   if host_id is None:
     host_id = get_backend(backend).host_id()
+  if host_id not in host_ids():
+    raise ValueError(f"Unknown host_id {host_id}")
   return [d for d in devices(backend) if d.host_id == host_id]
 
 
-def host_id(backend=None):
+def host_id(backend: str = None):
   """Returns the integer host ID of this host.
 
   On most platforms, this will always be 0. This will vary on multi-host
@@ -225,12 +242,12 @@ def host_id(backend=None):
   return get_backend(backend).host_id()
 
 
-def host_ids(backend=None):
+def host_ids(backend: str = None):
   """Returns a sorted list of all host IDs."""
   return sorted(list(set(d.host_id for d in devices(backend))))
 
 
-def host_count(backend=None):
+def host_count(backend: str = None):
   """Returns the number of hosts."""
   return len(host_ids(backend))
 


Test patch:
diff --git a/tests/xla_bridge_test.py b/tests/xla_bridge_test.py
--- a/tests/xla_bridge_test.py
+++ b/tests/xla_bridge_test.py
@@ -35,7 +35,7 @@ def test_set_device_assignment_with_partition(self):
                                   "0 2 \nComputation 1: 1 3 \n")
     self.assertEqual(compile_options.device_assignment.__repr__(),
                      expected_device_assignment)
-    
+
   def test_parameter_replication_default(self):
     c = xb.make_computation_builder("test")
     param = xb.parameter(c, 0, xc.Shape.array_shape(xc.PrimitiveType.F32, ()))
@@ -48,6 +48,13 @@ def test_parameter_replication(self):
     built_c = c.Build()
     assert "parameter_replication={false}" in built_c.as_hlo_text()
 
+  def test_local_devices(self):
+    self.assertNotEmpty(xb.local_devices())
+    with self.assertRaisesRegex(ValueError, "Unknown host_id 100"):
+      xb.local_devices(100)
+    with self.assertRaisesRegex(RuntimeError, "Unknown backend foo"):
+      xb.local_devices(backend="foo")
+
 
 if __name__ == "__main__":
   absltest.main()