Support XLA_CPU/GPU devices in eager runtime

Compiling a defun can be done with:

  @function.defun
  def f(x):
    return x + 1

  with tf.device("XLA_GPU"):
    f(5)

PiperOrigin-RevId: 224560868