from tensorflow.python.ops import state_ops


        lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype)
        mu_t = math_ops.cast(self._mu_t, var.dtype.base_dtype)
        vstar = self.get_slot(var, "vstar")
        gold = self.get_slot(var, "gold")  # glod is not sparse

        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():
