Initializes the TPUStrategy object. Args: tpu_cluster_resolver: A tf.distribute.cluster_resolver.TPUClusterResolver, which provides information about the TPU cluster. device_assignment: Optional `tf.tpu.experimental.DeviceAssignment` to specify the placement of r
(self,
tpu_cluster_resolver=None,
device_assignment=None)
| 86 | """TPU distribution strategy implementation.""" |
| 87 | |
| 88 | def __init__(self, |
| 89 | tpu_cluster_resolver=None, |
| 90 | device_assignment=None): |
| 91 | """Initializes the TPUStrategy object. |
| 92 | |
| 93 | Args: |
| 94 | tpu_cluster_resolver: A tf.distribute.cluster_resolver.TPUClusterResolver, |
| 95 | which provides information about the TPU cluster. |
| 96 | device_assignment: Optional `tf.tpu.experimental.DeviceAssignment` to |
| 97 | specify the placement of replicas on the TPU cluster. Currently only |
| 98 | supports the usecase of using a single core within a TPU cluster. |
| 99 | """ |
| 100 | super(TPUStrategy, self).__init__(TPUExtended( |
| 101 | self, tpu_cluster_resolver, device_assignment=device_assignment)) |
| 102 | |
| 103 | # TODO(cjfj): Modify `_call_for_each_replica` in `TPUExtended` such that this |
| 104 | # can use the default implementation. |
no test coverage detected