<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def calc_variance(grad_dict, num_batches, param_names):
 """Calculates the variance 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 param_names: str parameter name in the module Returns grad_dict: dict dictionary with new keys mapping to gradients variance """

    for i in range(len(param_names)):
        diff_sqr = mx.ndarray.square(mx.nd.subtract(grad_dict[param_names[i]],
                                                    grad_dict[str.format(param_names[i]+"_expectation")]))
        grad_dict[str.format(param_names[i] + "_variance")] = mx.ndarray.sum(diff_sqr, axis=0) / num_batches