def initialize_system(embedding_config=None, job=None): """Initializes a distributed TPU system for use with TensorFlow. Args: embedding_config: If not None, an `EmbeddingLayerConfiguration` proto describing the desired configuration of the hardware embedding lookup tables. If embedding_config is None, no hardware embeddings can be used. job: The job (the XXX in TensorFlow device specification /job:XXX) that contains the TPU devices that will be initialized. If job=None it is assumed there is only one job in the TensorFlow flock, and an error will be returned if this assumption does not hold. Returns: A serialized `TopologyProto` that describes the TPU system. Note: the topology must be evaluated using `Session.run` before it can be used. """ config_string = ("" if embedding_config is None else embedding_config.SerializeToString()) with ops.device(_tpu_system_device_name(job)): return tpu_ops.configure_distributed_tpu(embedding_config=config_string)
def initialize_system(embedding_config=None, job=None): """Initializes a distributed TPU system for use with TensorFlow. Args: embedding_config: If not None, a `TPUEmbeddingConfiguration` proto describing the desired configuration of the hardware embedding lookup tables. If embedding_config is None, no hardware embeddings can be used. job: The job (the XXX in TensorFlow device specification /job:XXX) that contains the TPU devices that will be initialized. If job=None it is assumed there is only one job in the TensorFlow flock, and an error will be returned if this assumption does not hold. Returns: A serialized `TopologyProto` that describes the TPU system. Note: the topology must be evaluated using `Session.run` before it can be used. """ config_string = ("" if embedding_config is None else embedding_config.SerializeToString()) with ops.device(_tpu_system_device_name(job)): return tpu_ops.configure_distributed_tpu(embedding_config=config_string)
def initialize_system(embedding_config=None, job=None): """Initializes a distributed TPU system for use with TensorFlow. Args: embedding_config: If not None, an EmbeddingLayerConfiguration proto describing the desired configuration of the hardware embedding lookup tables. If embedding_config is None, no hardware embeddings can be used. job: The job (the XXX in TensorFlow device specification /job:XXX) that contains the TPU devices that will be initialized. If job=None it is assumed there is only one job in the TensorFlow flock, and an error will be returned if this assumption does not hold. Returns: Op which, when executed, will initialize the system. """ if job is None: device_name = "/device:TPU_SYSTEM:0" else: device_name = "/job:%s/device:TPU_SYSTEM:0" % job config_string = ("" if embedding_config is None else embedding_config.SerializeToString()) with ops.device(device_name): init_distributed_tpu = tpu_ops.configure_distributed_tpu( embedding_config=config_string) return init_distributed_tpu
def initialize_system(embedding_config=None, job=None): """Initializes a distributed TPU system for use with TensorFlow. Args: embedding_config: If not None, an EmbeddingLayerConfiguration proto describing the desired configuration of the hardware embedding lookup tables. If embedding_config is None, no hardware embeddings can be used. job: The job (the XXX in TensorFlow device specification /job:XXX) that contains the TPU devices that will be initialized. If job=None it is assumed there is only one job in the TensorFlow flock, and an error will be returned if this assumption does not hold. Returns: Op which, when executed, will initialize the system. """ if job is None: device_name = "/device:TPU_SYSTEM:0" else: device_name = "/job:%s/device:TPU_SYSTEM:0" % job config_string = ("" if embedding_config is None else embedding_config.SerializeToString()) with ops.device(device_name): init_distributed_tpu = tpu_ops.configure_distributed_tpu( embedding_config=config_string) return init_distributed_tpu