Adding new Jax Models.
Is there any board to track that how many current models have jax implementation ?
I would like to contribute to add jax implementation for the remaining ones, which model I can take to start ?  
Hi there, thank you for your interest in JAX.

We plan to add as many models as possible in JAX/Flax. Right now we are working on improving the JAX support in the lib, better JAX/Flax tests, generation, cookie-cutter templates etc so that it'll become easier to add more models faster.

Please stay tuned, we'll soon share more details :)
