Generate the implementation of DistilledSGLD

def DistilledSGLD(teacher_sym, student_sym,
                  teacher_data_inputs, student_data_inputs,
                  X, Y, X_test, Y_test, total_iter_num,
                  teacher_learning_rate, student_learning_rate,
                  teacher_lr_scheduler=None, student_lr_scheduler=None,
                  student_optimizing_algorithm='sgd',
                  teacher_grad_f=None, student_grad_f=None,
                  teacher_prior_precision=1, student_prior_precision=0.001,
                  perturb_deviation=0.001,
                  student_initializer=None,
                  teacher_initializer=None,
                  minibatch_size=100,
                  task='classification',
                  dev=mx.gpu()):
    """Generate the implementation of DistilledSGLD"""
    teacher_exe, teacher_params, teacher_params_grad, _ = \
        get_executor(teacher_sym, dev, teacher_data_inputs, teacher_initializer)
    student_exe, student_params, student_params_grad, _ = \
        get_executor(student_sym, dev, student_data_inputs, student_initializer)
    if teacher_grad_f is None:
        teacher_label_key = list(set(teacher_data_inputs.keys()) - set(['data']))[0]
    if student_grad_f is None:
        student_label_key = list(set(student_data_inputs.keys()) - set(['data']))[0]
    teacher_optimizer = mx.optimizer.create('sgld',
                                            learning_rate=teacher_learning_rate,
                                            rescale_grad=X.shape[0] / float(minibatch_size),
                                            lr_scheduler=teacher_lr_scheduler,
                                            wd=teacher_prior_precision)
    student_optimizer = mx.optimizer.create(student_optimizing_algorithm,
                                            learning_rate=student_learning_rate,
                                            rescale_grad=1.0 / float(minibatch_size),
                                            lr_scheduler=student_lr_scheduler,
                                            wd=student_prior_precision)
    teacher_updater = mx.optimizer.get_updater(teacher_optimizer)
    student_updater = mx.optimizer.get_updater(student_optimizer)
    start = time.time()
    for i in range(total_iter_num):
        # 1.1 Draw random minibatch
        indices = numpy.random.randint(X.shape[0], size=minibatch_size)
        X_batch = X[indices]
        Y_batch = Y[indices]

        # 1.2 Update teacher
        teacher_exe.arg_dict['data'][:] = X_batch
        if teacher_grad_f is None:
            teacher_exe.arg_dict[teacher_label_key][:] = Y_batch
            teacher_exe.forward(is_train=True)
            teacher_exe.backward()
        else:
            teacher_exe.forward(is_train=True)
            teacher_exe.backward(
                teacher_grad_f(teacher_exe.outputs, nd.array(Y_batch, ctx=dev)))

        for k in teacher_params:
            teacher_updater(k, teacher_params_grad[k], teacher_params[k])

        # 2.1 Draw random minibatch and do random perturbation
        if task == 'classification':
            indices = numpy.random.randint(X.shape[0], size=minibatch_size)
            X_student_batch = X[indices] + numpy.random.normal(0,
                                                               perturb_deviation,
                                                               X_batch.shape).astype('float32')
        else:
            X_student_batch = mx.random.uniform(-6, 6, X_batch.shape, mx.cpu())

        # 2.2 Get teacher predictions
        teacher_exe.arg_dict['data'][:] = X_student_batch
        teacher_exe.forward(is_train=False)
        teacher_pred = teacher_exe.outputs[0]
        teacher_pred.wait_to_read()

        # 2.3 Update student
        student_exe.arg_dict['data'][:] = X_student_batch
        if student_grad_f is None:
            student_exe.arg_dict[student_label_key][:] = teacher_pred
            student_exe.forward(is_train=True)
            student_exe.backward()
        else:
            student_exe.forward(is_train=True)
            student_exe.backward(student_grad_f(student_exe.outputs, teacher_pred))
        for k in student_params:
            student_updater(k, student_params_grad[k], student_params[k])

        if (i + 1) % 2000 == 0:
            end = time.time()
            if task == 'classification':
                print("Current Iter Num: %d" % (i + 1), "Time Spent: %f" % (end - start))
                test_correct, test_total, test_acc = \
                    sample_test_acc(student_exe, X=X_test, Y=Y_test, label_num=10,
                                    minibatch_size=minibatch_size)
                train_correct, train_total, train_acc = \
                    sample_test_acc(student_exe, X=X, Y=Y, label_num=10,
                                    minibatch_size=minibatch_size)
                teacher_test_correct, teacher_test_total, teacher_test_acc = \
                    sample_test_acc(teacher_exe, X=X_test, Y=Y_test, label_num=10,
                                    minibatch_size=minibatch_size)
                teacher_train_correct, teacher_train_total, teacher_train_acc = \
                    sample_test_acc(teacher_exe, X=X, Y=Y, label_num=10,
                                    minibatch_size=minibatch_size)
                print("Student: Test ACC %d/%d=%f, Train ACC %d/%d=%f" % (test_correct, test_total,
                                                                          test_acc, train_correct,
                                                                          train_total, train_acc))
                print("Teacher: Test ACC %d/%d=%f, Train ACC %d/%d=%f" \
                      % (teacher_test_correct, teacher_test_total, teacher_test_acc,
                         teacher_train_correct, teacher_train_total, teacher_train_acc))
            else:
                print("Current Iter Num: %d" % (i + 1), "Time Spent: %f" % (end - start), "MSE:",
                      sample_test_regression(exe=student_exe, X=X_test, Y=Y_test,
                                             minibatch_size=minibatch_size,
                                             save_path='regression_DSGLD.txt'))
            start = time.time()

    return student_exe, student_params, student_params_grad