o
    i e                     @   sf   d Z ddlm  mZ ddlmZ ejZeddd eddd ed	dd
 dd Z	dd Z
dS )z.Utility functions for tests using TPUStrategy.    N)flagstpu zName of TPU to connect to.projectzName of GCP project with TPU.zonezName of GCP zone with TPU.c                  C   s   t jjjtjtjtjd} | S )N)r   r   r   )tf
distributecluster_resolverTPUClusterResolverFLAGSr   r   r   resolver r   [/var/www/myenv/lib/python3.10/site-packages/keras/src/distribute/tpu_strategy_test_utils.pyget_tpu_cluster_resolver   s   r   c                  C   s.   t  } tj|  tjj|  tjj| S )N)	r   r   configexperimental_connect_to_clusterr   experimentalinitialize_tpu_systemr   TPUStrategyr   r   r   r   get_tpu_strategy#   s   r   )__doc__tensorflow.compat.v2compatv2r   abslr   r   DEFINE_stringr   r   r   r   r   r   <module>   s   	