Split TowerContext to TrainTowerContext and PredictTowerContext to support more informative contexts in the future