Optimizer helper functions

[si-you]

SGD or Adagrad with a square root decay learning rate schedule can help training. Since TensorFlow lattice estimators accepts arbitrary callable as an optimizer, so we can use learning rate scheduling, but it is not easy to configure for a beginner. This pull requests include some helper functions and tests to illustrate how to use a custom learning rate schedule.

Also this pull requests include modify keep_dims argument to keepdims in linear model construction since tensorflow is deprecating a former argument.

[mmilanifard]

Discussed this internally. Although this is useful, it is not directly linked to this project. We should at least push this under lib directory.

[si-you]

Hi Mahdi. Yes it is not specific to our estimator classes, rather this function can be used for any estimator classes, so it would be handy to check in this to tensorfow.estimator directly but not sure how early we can do that :).

Do you mean lib like tensorflow_lattice/python/lib?

[mmilanifard]

I think having it in tensorflow repo itself if better. I don't think lattice estimators is a good place for this.
By lib I meant python/lib in this repo. But that is not ideal either.

[si-you]

How about writing a new example that illustrates how to use learning rate scheduling and check that in under tensorflow_lattice/examples/ ? As I said, it is somewhat non-trivial to use learning rate scheduling in estimator :(.

[mmilanifard]

That sounds good. Having clear examples outside the test is often a good idea.