parameters: (
          server_rref,
          received_batch_number,
          param_loc,
          gradient
      )
return_statement: return fut
docstring: r"""
          An asynchronous function that will average gradients
          sent from trainers.
          Args:
              server_rref (RRef): remote reference to the server
              received_batch_number (int): batch number sent by
                  the trainer
              param_loc (int): bucket location sent by the trainer
                  containing the gradient
              gradient (torch.Tensor or list): tensor sent by the trainer
docstring_summary: r"""
          An asynchronous function that will average gradients
          sent from trainers.
          Args:
              server_rref (RRef): remote reference to the server
              received_batch_number (int): batch number sent by
                  the trainer
              param_loc (int): bucket location sent by the trainer
                  containing the gradient
              gradient (torch.Tensor or list): tensor sent by the trainer
docstring_tokens: r An asynchronous function that will average gradients sent from trainers . Args : server_rref ( RRef ) : remote reference to the server received_batch_number ( int ) : batch number sent by the trainer param_loc ( int ) : bucket location sent by the trainer containing the gradient gradient ( torch . Tensor or list ) : tensor sent by the trainer
function: def average_gradient(
          server_rref,
          received_batch_number,
          param_loc,
          gradient
      ):
          r"""
          An asynchronous function that will average gradients
          sent from trainers.
          Args:
              server_rref (RRef): remote reference to the server
              received_batch_number (int): batch number sent by
                  the trainer
              param_loc (int): bucket location sent by the trainer
                  containing the gradient
              gradient (torch.Tensor or list): tensor sent by the trainer
          """
          self = server_rref.local_value()
          if type(gradient) is list:
              gradient = sparse_rpc_format_to_tensor(gradient)
          gradient = gradient.cuda(self.rank)
          fut = torch.futures.Future()
          with self.lock:
              if self.batch_number < received_batch_number:
                  self.batch_number = received_batch_number
                  self.clear_batch_state()
              self.process_gradient(gradient, param_loc)
              if param_loc not in self.futures:
                  self.futures[param_loc] = []
              self.futures[param_loc].append(fut)
              if len(self.futures[param_loc]) == self.trainer_count:
                  self.record_straggler_end(self.param_key(param_loc))
                  param_loc_avg = self.average(param_loc)
                  if not self.use_cuda_rpc:
                      param_loc_avg = param_loc_avg.cpu()
                  if param_loc_avg.is_sparse:
                      param_loc_avg = sparse_tensor_to_rpc_format(param_loc_avg)
                  for cur_fut in self.futures[param_loc]:
                      cur_fut.set_result(param_loc_avg)
                  self.record_batch_end(self.param_key(param_loc))
          return fut
function_tokens: def average_gradient ( server_rref , received_batch_number , param_loc , gradient ) : self = server_rref . local_value ( ) if type ( gradient ) is list : gradient = sparse_rpc_format_to_tensor ( gradient ) gradient = gradient . cuda ( self . rank ) fut = torch . futures . Future ( ) with self . lock : if self . batch_number < received_batch_number : self . batch_number = received_batch_number self . clear_batch_state ( ) self . process_gradient ( gradient , param_loc ) if param_loc not in self . futures : self . futures [ param_loc ] = [ ] self . futures [ param_loc ] . append ( fut ) if len ( self . futures [ param_loc ] ) == self . trainer_count : self . record_straggler_end ( self . param_key ( param_loc ) ) param_loc_avg = self . average ( param_loc ) if not self . use_cuda_rpc : param_loc_avg = param_loc_avg . cpu ( ) if param_loc_avg . is_sparse : param_loc_avg = sparse_tensor_to_rpc_format ( param_loc_avg ) for cur_fut in self . futures [ param_loc ] : cur_fut . set_result ( param_loc_avg ) self . record_batch_end ( self . param_key ( param_loc ) ) return fut