Precise control of Device with default_device

[joeryjoery]

I was trying to get jax.default_device working as I want to have more control on how to place computation, but it seems that jax has no attribute to this context manager? Even though it would seem it that has to, right? see #9118

import jax

jax.default_device()
>> Traceback (most recent call last):
    File "<stdin>", line 1, in <module>
  AttributeError: module 'jax' has no attribute 'default_device'

https://jax.readthedocs.io/en/latest/_autosummary/jax.default_device.html

I was trying to hack around this with setting jax.jit(..., device=device) on my desired function, however because I'm using this iteratively, I would get loads of back and forth copying from the global/ default device to the computation device (causing my computer to freeze up and hog RAM that would make even Chrome jealous). I'm aware of the global device setting with e.g., CUDA_VISIBLE_DEVICE etc., but a Tensorflow session like context manager would be optimal.

Does anyone know if I missed something? I'm using Jax version 0.3.5.

[wookayin]

jax.default_device was added in #9118, available in jax 0.3.14 or higher.