gcloud alpha compute tpus tpu-vm create $TPU_NAME \
        --zone=$TPU_ZONE \
        --accelerator-type=$ACCELERATOR_TYPE \
        --version=$VERSION

def get_tpu_strategy() -> tf.distribute.Strategy:
    """Get tensorflow distribution strategy for TPU

    Returns
    -------
    tf.distribute.Strategy
        Model distribution strategy for TPU

    """
    resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
        tpu=os.environ["TPU_NAME"],
        zone=os.environ["TPU_ZONE"],
        project=os.environ["GCP_PROJECT"],
    )
    logger.info(f"All TPU devices: {resolver.cluster_spec().as_dict()['worker']}")
    tf.config.experimental_connect_to_cluster(resolver)
    logger.info("Connected to TPUCluster")
    tf.tpu.experimental.initialize_tpu_system(resolver)
    strategy = tf.distribute.TPUStrategy(resolver)
    logger.info(f"Number of TPU accelerators: {strategy.num_replicas_in_sync}")
    strategy.num_workers = len(
        strategy._extended._get_input_workers(options=None)._worker_device_pairs
    )
    return strategy