Update structure of MNIST example.

[copybara-service]

Update structure of MNIST example.

• mnist_main.py -> main.py
• mnist_lib.py -> train.py

The main.py shouldn't be used in any test.

[codecov-io]

Codecov Report

Merging #634 (e165074) into master (80422e5) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##           master     #634   +/-   ##
=======================================
  Coverage   80.59%   80.59%           
=======================================
  Files          56       56           
  Lines        4241     4241           
=======================================
  Hits         3418     3418           
  Misses        823      823           

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 80422e5...e165074. Read the comment docs.

[andsteing]

Quoting from #231 :

• To make it easier to test, maintain and reuse code structure the code as follows:

  • {example}_lib.py contains classes and methods for training and evaluating the model.
  • {example}_main.py contains the flags and calls a method from {example}_lib.py to run the training loop. This should be the only file defining flags!
  • {example}_lib_test.py for test cases of the training code. At a minimum the test should run a single training step but more fine grained unit tests are a bonus. You can use tfds.testing.mock_data to avoid real data from disk.

I'm fine with any structure, but it should be the same for every example. If you think the structure proposed in this PR is superior, we should update #231 and make sure all of the examples follow the updated structure.

[Marvin182]

I agree. I updated #231.

[andsteing]

Please also update the local imports in the Colab:

https://colab.sandbox.google.com/github/google/flax/blob/test_342404042/linen_examples/mnist/mnist.ipynb#scrollTo=6Y1ru2Ovp4ZI

[Marvin182]

done