tensorflow: is there a way to specify XLA_GPU with tensorflow?

Question

following code is used to specify device on which tf node is running on

with tf.device('/gpu:0'):
  a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')

i have already known: this post, tensorflow doc and xla demo

what i want to know is:

is there a way to specify XLA_GPU as the device on which tf node is running on

with tf.device('/XLA_GPU:0'):
  a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')

executing the code above gives ValueError: Unknown attribute: 'device' in '/job:localhost/replica:0/task:0/device:/XLA_GPU:0'

this is 100% reproducible on google colab.

Answer 1

try 'device:XLA_GPU:0'

with tf.device('device:XLA_GPU:0'):
  a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')