Updating linen_examples/vae

[andsteing]

This issue tracks updates the vae example to follow practices outlined in #231.

• [x] Port to linen API - once ported all subsequent changes should be done in linen_examples/vae
• [ ] Update README.md
• [x] Add requirements.txt
• [ ] Update file structure
• [ ] Use ml_collections.ConfigDict ?
• [ ] Add benchmark test ?
• [ ] Add unit test for training/eval step
• [ ] Add Colab
• [ ] Adhere to Google Python style
• [ ] Add mypy annotations
• [ ] Shorten/beautify training loop (consider using clu for this)

[andsteing]

Note that "file structure" in #231 got changed (as part of #634):

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

  • main.py contains the flags and calls a method from train.py to run the training loop. This should be the only file defining flags! This can be almost identical for all examples, please copy from linen_examples/wmt/main.py.
  • train.py contains classes and methods for training and evaluating the model.
  • train_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.

[marcvanzee]

Unassigning myself and adding "pull requests welcome"

[canyon289]

Mind if I take a run at this? I might not be able to do everything in the checklist but can certainly move this along

[andsteing]

Yes, please have a run.

Happy to answer questions and review PRs.

[canyon289]

Thanks! Feel free to assign this to me, I'll likely open up multiple smaller PRs for subpieces rather than one open one. If that doesn't sound like a good plan please let me know!

[SauravMaheshkar]

We should reopen the issue and mark some of the tasks as done

[SauravMaheshkar]

Even mypy and Google Style Guide can be checked IMO. Gentle ping @levskaya

[SauravMaheshkar]

Updated Tracker

• [x] Port to linen API
• [x] Update README.md
• [x] Add requirements.txt
• [x] Update file structure
• [x] Use ml_collections.ConfigDict
• [ ] Add benchmark test
• [ ] Add unit test for training/eval step
• [ ] Add Colab
• [x] Adhere to Google Python style
• [x] Add mypy annotations
• [ ] Shorten/beautify training loop (consider using clu for this)

[shimejii]

Request for permission to modify

The current vae example does not work, may I make minor modifications?

Modifications

Two modifications are as follows:

1. modify requirements.txt
   • add clu==0.0.6
   • upgrade flax==0.6.9 To flax==0.7.4
2. modify README.md
   • remove "--workdir=/tmp/mnist" from a execution command

Reasons

• flax==0.7.4

  • Update handling of typed PRNG keys in version 0.7.3

    • flax_v0.7.3

  • flax -0.7.3 is yanked

    • pypi_flax_v0.7.3

• remove "--workdir=/tmp/mnist"

  • commit hash db7b1762
    • workdir flag definition is removed in main.py

[andsteing]

Yes, please go ahead and create a PR to fix the example.

[shimejii]

Thank you very much. I have created a PR, please check it.