Flax Masked Language Modeling training example
Include a training example running with Flax/JAX framework. (cc @avital @marcvanzee)

TODOs:

- [x] Make the collator working with Numpy/JAX array 
- [x] Make sure the training actually works on larger scale
- [x] Make it possible to train from scratch
- [x] Support TPU (`bfloat16`)
- [ ] Support GPU amp (`float16`)
- [x] Improve overall UX