def _sync_weights_and_state_across_hosts(self):
    """Sync weights and state across all the hosts in the computation."""

    if logging.vlog_is_on(1):
      logging.debug(
          'Input training weights shape: %s',
          fastmath.nested_map(lambda x: x.shape,
                              self._model.weights))
      logging.debug('Input training weights: %s', self._model.weights)
      logging.debug('Input training state: %s', self._model.state)
      logging.debug('Input eval weights: %s', self._eval_model.weights)
      logging.debug('Input eval state: %s', self._eval_model.state)

    (self._model.weights, self._model.state,
     self._eval_model.weights, self._eval_model.state) = self._unreplicate(
         _make_weights_and_state_same_across_hosts(
             self._for_n_devices(
                 (self._model.weights, self._model.state,
                  self._eval_model.weights,
                  self._eval_model.state))))

    if logging.vlog_is_on(1):
      logging.debug(
          'Output training weights shape: %s',
          fastmath.nested_map(lambda x: x.shape, self._model.weights))
      logging.debug('Output training weights: %s', self._model.weights)
      logging.debug('Output training state: %s', self._model.state)
      logging.debug('Output eval weights: %s', self._eval_model.weights)
      logging.debug('Output eval state: %s', self._eval_model.state)
Sync weights and state across all the hosts in the computation.
trax/supervised/training.py