Simplify and test the training code

[srush]

Going through my journey of learning Flax, updated some of the code for training.

Changes:

• Factor out data and training into seperate files for debugging.
• Make sure that state is threaded through all of the training code so that it trains
• Change FeedForward and LSTM so they are used as part of seqmodel
• Heavy use of vmap and vectorize to ensure that none of the code has extra dimension or indexing.
• Added a main file so that training.py can be run directly.

I haven't updated the comments yet or literate aspects yet.

[siddk]

This makes sense to me - didn't follow the use of @partial(np.vectorize) at first, might be worth a comment.

I'm running into some relative import errors when I try to run the code in this PR, will fix, and push... ideally hold of on merging until then!