Repository: jax-ml/jax (python)

Issue:
Array API `default_device()` and `devices()` are incorrect
### Description

`jax.numpy.__array_namespace_info__().default_device()` returns None.
It should return a Device object (read https://github.com/data-apis/array-api/issues/835 for controversy on which device should be returned).

`jax.numpy.__array_namespace_info__().devices()` is an alias to `jax.devices()`.
While this is not _explicitly_ spelled out by the Array API (https://data-apis.org/array-api/latest/API_specification/generated/array_api.info.devices.html#devices), I believe it is misleading as it lacks the `backend` optional parameter that `jax.devices` offers. I think it should return a list of _all_ devices available at the moment of calling it; e.g. when CUDA is available, it should return `[CpuDevice(id=0), CudaDevice(id=0)]`.

### System info (python version, jaxlib version, accelerator, etc.)

jax-0.5.2
jaxlib-0.5.2-cuda126

Pull request description:
[array api] return all devices in devices()
Fixes #27606

Patch:
diff --git a/jax/_src/numpy/array_api_metadata.py b/jax/_src/numpy/array_api_metadata.py
index d8d2c2d1a..5267e5121 100644
--- a/jax/_src/numpy/array_api_metadata.py
+++ b/jax/_src/numpy/array_api_metadata.py
@@ -24,7 +24,9 @@ from types import ModuleType
 import jax
 from jax._src.sharding import Sharding
 from jax._src.lib import xla_client as xc
-from jax._src import dtypes as _dtypes, config
+from jax._src import config
+from jax._src import dtypes as _dtypes
+from jax._src import xla_bridge as xb
 
 
 __array_api_version__ = '2023.12'
@@ -73,7 +75,10 @@ class ArrayNamespaceInfo:
     return None
 
   def devices(self):
-    return jax.devices()
+    out = [None]  # None indicates "uncommitted"
+    for backend in xb.backends():
+        out.extend(jax.devices(backend))
+    return out
 
   def capabilities(self):
     return self._capabilities


Test patch:
diff --git a/tests/array_api_test.py b/tests/array_api_test.py
index d509fe78c..8e4ba275f 100644
--- a/tests/array_api_test.py
+++ b/tests/array_api_test.py
@@ -26,6 +26,7 @@ import jax
 import jax.numpy as jnp
 from jax._src import config, test_util as jtu
 from jax._src.dtypes import _default_types, canonicalize_dtype
+from jax._src import xla_bridge as xb
 
 ARRAY_API_NAMESPACE = jnp
 
@@ -283,7 +284,10 @@ class ArrayAPIInspectionUtilsTest(jtu.JaxTestCase):
     assert self.info.default_device() is None
 
   def test_devices_info(self):
-    assert self.info.devices() == jax.devices()
+    devices = set(self.info.devices())
+    assert None in devices
+    for backend in xb.backends():
+      assert devices.issuperset(jax.devices(backend))
 
   def test_default_dtypes_info(self):
     _default_dtypes = {