tf.experimental.dtensor.initialize_tpu_system

Initialize the TPU devices.

tf.experimental.dtensor.initialize_tpu_system(
    enable_coordination_service=False
)

This functions performs additional TPU related initialization after calling dtensor.initialize_multi_client to initialize multi-client DTensor. Refer to dtensor.initialize_multi_client for relevant environment variables that controls the initialization of multi-client DTensor.

enable_coordination_service If true, enable distributed coordination service to make sure that workers know the devices on each other, a prerequisite for data transfer through cross-worker rendezvous.

RuntimeError If running inside a tf.function.
NotFoundError If no TPU devices found in eager mode.