tf.experimental.dtensor.num_global_devices

Returns the number of devices of device_type in this DTensor cluster.