Add unit tests (smoke tests + reproducibility tests) for the JFT training scripts.

[dusenberrymw]

That way we can catch errors and any changes that inadvertently affect reproducibility as we refactor and make changes. Especially helpful for large pre-training models so that we can spend more time on experiments than debugging.