def _init_strategy(self):
    """Initialize the distribution strategy (e.g. TPU/GPU/Mirrored)."""
    if self._strategy is None:
      if self._tpu is not None:
        resolver = distribute_utils.tpu_initialize(self._tpu)
        self._strategy = tf.distribute.experimental.TPUStrategy(resolver)
      elif (
          self._distribution_strategy is None
          or self._distribution_strategy == 'default'
      ):
        self._strategy = tf.distribute.get_strategy()
      elif self._distribution_strategy == 'cpu':
        self._strategy = tf.distribute.OneDeviceStrategy('/device:cpu:0')
      else:
        if self._distribution_strategy == 'mirrored':
          self._strategy = tf.distribute.MirroredStrategy()
        else:
          raise ValueError(
              f'Invalid distribution strategy="{self._distribution_strategy}"')
Initialize the distribution strategy (e.g. TPU/GPU/Mirrored).
language/serene/training.py