[rhdunn]

It's useful to be able to run models/code on the CPU or split between CPU and GPU, especially for models that cannot fit into the GPU's VRAM, if you are running multiple models, or if you have training data that you need to move between CPU and GPU.

[earthnail]

All of that would be possible with the changes parent poster proposes. See the gotchas section in JAX, which is exactly these limitations:

https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_J...