tensorflow: is there a way to specify XLA_GPU with tensorflow? The following code is used to specify device on which tf node is running

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 any way to specify XLA_GPU as the device on which tf node is running?

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.
Check that the device is there first by examining the device names in the output from the following:

from tensorflow.python.client import device_lib
device_lib.list_local_devices()


Once you've verified that 'XLA_GPU:0' is a device on your system use something like 

with tf.device('/device:XLA_GPU:0'):


Obviously you can change the names.