Refactor trainer class to expose the train_with_lr().
Refactor trainer class to expose the train_with_lr().

This is particularly useful to get the main train_step functions out of trainer.

PiperOrigin-RevId: 437805544

google-research/t5x