def METHOD_NAME(self, grad, var, indices, apply_state=None):
    if "apply_state" in self._optimizer._sparse_apply_args:
        train_op = self._optimizer.METHOD_NAME(
            grad, var, indices, apply_state=apply_state
        )
    else:
        train_op = self._optimizer.METHOD_NAME(grad, var, indices)
    average_op = self._apply_average_op(train_op, var, apply_state)
    return tf.group(train_op, average_op)