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

Question

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.

Answer 1

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.