tf.experimental.dtensor.device_name

Returns the singleton DTensor device's name.

This function can be used in the following way:

import tensorflow as tf

with tf.device(dtensor.device_name()):
  # ...