flax

872d50b8 - [flax:examples:nlp_seq] Create a main.py file to run tests with config files to match other examples. #jax-fixit