[utils]

Utils for jax backend support code.

GPU/CPU backing device management utils: * select_device: Select a backing device to use based on inputs and availability.