Updates the ensembling HOWTO to Optax.
This commit also bases the code on the more recent version of the MNIST example and
implements the actual ensembling (before this commit, the individual
members' predictions weren't actually ensembled, but instead the accuracy
was incorrectly averaged from individual members' accuracies).
TESTED_MANUALLY:
https://colab.research.google.com/drive/1Hgc2yYdMbrVNEELqZtF7vr6aIYF3WknS
