xla

f5b33c5d - training UX: automatic generating make_train_step (#8495)