Flax from_pretrained: clean up mismatched_keys.
#630
Merged
from_pretrained: clean up mismatched_keys.
#630