How to get the current default device set by jax.default_device

[muupan]

Is there a straightforward way to get the current default device of jax? I could not find such a function in the API documentation.

I tried jax.devices()[0] but found it does not reflect the effect of jax.default_device.

print(jax.devices()[0])

with jax.default_device(jax.devices("cpu")[0]):
    print(jax.devices()[0])

gpu:0
gpu:0

As a workaround, I found that making a new array and checking its device works. But this is not straightforward and I feel there should be a better way.

print(jnp.empty(0).device())

with jax.default_device(jax.devices("cpu")[0]):
    print(jnp.empty(0).device())

gpu:0
TFRT_CPU_0

[mattjj]

Thanks for the question!

I think the current (undocumented) answer is that you can check the value of jax.config.jax_default_device, which is an Optional[Device], and if it's None then use jax.local_devices()[0]. That is:

import jax

def get_default_device():
  return jax.config.jax_default_device or jax.local_devices()[0]

What do you think?

[muupan]

Thanks, your method worked.

def get_default_device():
  return jax.config.jax_default_device or jax.local_devices()[0]

print(get_default_device())

with jax.default_device(jax.devices("cpu")[0]):
    print(get_default_device())
    
print(get_default_device())

gpu:0
TFRT_CPU_0
gpu:0