code: def calc_expectation(grad_dict, num_batches):
      
      for key in grad_dict.keys():
          grad_dict[str.format(key+"_expectation")] = mx.ndarray.sum(grad_dict[key], axis=0) / num_batches
  
      return grad_dict
docs: Calculates the expectation of the gradients per epoch for each parameter w.r.t number of batches
  
      Parameters
      ----------
      grad_dict: dict
          dictionary that maps parameter name to gradients in the mod executor group
      num_batches: int
          number of batches
  
      Returns
      ----------
      grad_dict: dict
          dictionary with new keys mapping to gradients expectations