以下代码用于指定运行 tf 节点的设备

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')

我已经知道： 这篇文章、tensorflow doc和xla demo

我想知道的是：

有没有办法将 XLA_GPU 指定为运行 tf 节点的设备？

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')

执行上面的代码给出ValueError: Unknown attribute: 'device' in '/job:localhost/replica:0/task:0/device:/XLA_GPU:0'

这在 google colab 上是 100% 可重现的。