Array API `default_device()` and `devices()` are incorrect

[crusaderky]

Description

jax.numpy.__array_namespace_info__().default_device() returns None.
It should return a Device object (read data-apis/array-api#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

[jakevdp]

Hi - thanks for the report! This is working as intended. The meaning of None here is "uncommitted", which is JAX's default placement when none is specified:

jax/jax/_src/numpy/array_api_metadata.py

Lines 70 to 72 in 95497ca

	def default_device(self):
	# By default JAX arrays are uncommitted (device=None), meaning that
	# JAX is free to choose the most efficient device placement.

There was a long discussion around this when we initially implemented it, but I can't find a reference to it at the moment.

It should return a Device object

The array API does not restrict the type of the object that should be returned, only that the object should be valid when passed to the device argument of and array API function. The current default_device return value of None meets these requirements.

[crusaderky]

Could you comment on the second half of the issue?
jax.numpy.__array_namespace_info__().devices() doesn't return all available devices.

[jakevdp]

Ah, thanks I missed that. jax.numpy.__array_namespace_info__().devices() is just an alias to jax.devices(). Do you think it should have different behavior?

[crusaderky]

Yes; because jax.devices offers a backend parameter and the Array API variant doesn't, it should offer all devices on all backends:

def devices():
    out = []
    for backend in backends():
        out.extend(jax.devices(backend)
    return out

Otherwise, it is impossible for Array API agnostic code to discover that the library has multiple backends available. This impacts tests such as scipy/scipy#22756

[jakevdp]

Sure, that sounds reasonable.

Do you think we should also include None in the devices() output, to indicate that it is a device assignment value with special meaning?

[crusaderky]

I think it should be included if anything so that the output of default_device() is in the list returned by devices().

[jakevdp]

One other question: JAX arrays in general may be sharded across multiple devices, and so JAX allows any supported sharding specification to be passed to the device argument of array API functions. We can't enumerate all possible shardings (the number of valid specifications is more-or-less infinite) so we'd probably want to limit the output of devices() to just single devices – is that at all problematic? Does anything in the array API specification expect that arr.device in devices() should always return True?

[crusaderky]

Does anything in the array API specification expect that arr.device in devices() should always return True?

It doesn't, but I think it should be made explicit that one shouldn't expect it.