Flax: Jax Neural Network Library

Will be keeping snippets and concepts related to Flax here.

…