Weekly updated open sourced model implementations in Flax

[DarshanDeshpande]

JAX Models repository aims to provide open sourced JAX/Flax implementations for research papers originally without code or code written with frameworks other than JAX. The goal of this project is to make a collection of models, layers, activations and other utilities that are most commonly used for research

Some features of jax_models:

1. Pip installable
2. Weekly / fortnightly model additions
3. Comprehensive tests
4. Model architectures are cross checked by checking original implementations if available or otherwise miscellaneous aspects of the paper such as model parameters

The repository is currently in its pre-alpha stage with only around 6 model and layer implementations. If you are interested in this, you can contribute by either asking for specific paper implementations or by helping me train and open source weights of these models.

If you like my work then please consider starring the repository and giving feedback. Thank you!

Link: https://github.com/DarshanDeshpande/jax-models

[avital]

Thanks for sharing @DarshanDeshpande! Now that you've implemented these models, please let us know if there were challenges you hit, so that we can improve Flax going forward for others!

[DarshanDeshpande]

There's nothing as of now. Flax is really mature and flexible as it is.

If you could just direct me to some good benchmarks comparing Flax with other frameworks like TF or Pytorch, it would be very helpful. All the ones that I keep finding aren't comprehensive enough :(

[andsteing]

See also @matthias-wright's repository of pretrained models (mentioned in #1320):
https://github.com/matthias-wright/flaxmodels

When re-implementing models it might be particularly useful to provide code for loading PyTorch weights and adding tests verifying that both implementations produce identical results.

[matthias-wright]

Sounds good @marcvanzee! Are you thinking about something that looks like this ?

[marcvanzee]

Awesome! Yes, indeed, I was referring to a HOWTO, which we also sometimes call a pattern, tutorial, ...

We are still thinking how to best structure our guides, if you have ideas for how to improve this feel free to chime in as well.
In any case, for now a HOWTO seems like the right place.

[matthias-wright]

Okay cool, I will then submit a PR once I have something.

[andsteing]

This is now submitted as part of #1848
https://flax.readthedocs.io/en/latest/howtos/convert_pytorch_to_flax.html

Thanks for the nice PR, Matthias!

[DarshanDeshpande]

@andsteing @matthias-wright Do you think it's a good idea for the tutorial to also mention the approximations made by some layers? For example: LayerNorm or the default approximate=True as seen for gelu?