xla

ae042496 - Add linear scaling, warmup epochs and poly decay. correct few LARS steps and add feature to tune flags