import tensorflow as tf

        v_diff = state_ops.assign(vstar, mu_t * (var - vstar), use_locking=self._use_locking)

        with ops.control_dependencies([v_diff]):  # run v_diff operation before scatter_add
            scaled_grad = scatter_add(vstar, indices, grad)
        var_update = state_ops.assign_sub(var, lr_t * (scaled_grad + gold))

        return control_flow_ops.group(*[var_update, ])

    def _apply_sparse(self, grad, var):  # sparse grad (only for the shakespeare model)
        return self._apply_sparse_shared(
        grad.values, var, grad.indices, lambda x, i, v: state_ops.scatter_add(x, i, v))

    def set_params(self, cog, avg_gradient, client):
        with client.model.graph.as_default():
            all_vars = tf.trainable_variables()
            for variable, value in zip(all_vars, cog):
                vstar = self.get_slot(variable, "vstar")
                vstar.load(value, client.model.sess)
        
        # get old gradient
        _, gprev = client.get_grads()

        # Find g_t - F'(old)
        gdiff = [g1 - g2 for g1, g2 in zip(avg_gradient, gprev)]

        with client.model.graph.as_default():
            all_vars = tf.trainable_variables()
            for variable, grad in zip(all_vars, gdiff):
