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

def calc_expectation(grad_dict, num_batches):
    """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

    """
    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