changing api to have average_sgd and plain_sgd