Add fake enum-fields even during training in order to prevent pjit failures
Add fake enum-fields even during training in order to prevent pjit failures

Currently, when doing eval during training, we use the train input batch to get the partition specs used for pjit'ing the eval step. However, this does not work if the eval input batch has different data than what the train inputs have. This cl specifically matches the shapes between train and eval by adding fake enum fields to the train data to get consistent input partition specs to be used for eval. Stop gap until we update core training loop logic.

Users using pjit + enum in eval, should turn on `use_enumeration=True` in their train dataset if running train_and_eval with pjit.

PiperOrigin-RevId: 477823309

google/paxml,google/paxml