Remove boilerplate outside model, moving into Keras model for various baselines

In BatchEnsemble, we have to [add code like `tf.tile`](https://github.com/google/edward2/blob/82b7582ef99eeea014ebbec81fa9856f0cb3f60a/experimental/marginalization_mixup/batchensemble.py#L331) before passing inputs into the model, and splitting and average probs before passing into metrics.

This is fine in individual scripts, but once you start applying SavedModels for post-training pipelines, these lines have to be copied over (e.g., [`temperature_scaling.py`](https://github.com/google/edward2/blob/82b7582ef99eeea014ebbec81fa9856f0cb3f60a/experimental/marginalization_mixup/temperature_scaling.py) takes an arbitrary SavedModel as input; and Robustness Metrics). These scripts aim to be general-purpose, but specific baselines need further changes to their inputs/outputs. This makes the post-training pipelines not as general as they could be.