(self)
Sync weights and state across all the hosts in the computation.
Sync weights and state across all the hosts in the computation.
Sync weights and state across all the hosts in the computation .
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)
def _sync_weights_and_state_across_hosts ( self ) : 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 )