cloud_tpu_client_tensorflow_initialization_and_device_verification.py

1import tensorflow as tf
2from cloud_tpu_client import Client
3
4# Initialize the Cloud TPU client
5# If running on a TPU VM, no arguments are needed. 
6# Otherwise, provide the TPU name: Client(tpu='my-tpu-name')
7c = Client()
8
9# Configure the TPU runtime version
10c.configure_tpu_version('2.12.0', restart_type='ifNeeded')
11
12# Connect to the TPU cluster
13resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
14tf.config.experimental_connect_to_cluster(resolver)
15tf.tpu.experimental.initialize_tpu_system(resolver)
16
17# Verify the available devices
18devices = tf.config.list_logical_devices('TPU')
19print(f"Number of TPU devices: {len(devices)}")
20print("TPU Devices:", devices)