Implement Flux

[monatis]

There's a rough implementation by fellow GDEs, and they are interested in this project as well. It would be great if they migrate their work to this library also by implement the T5 part in pure JAX as well. It will be a great addition to this library.

[SauravMaheshkar]

Hey 👋🏼 , finally have some bandwidth to work on this. Let's split this into 2 separate issues ? I'm thinking I'll start with a T5 implementation and then we can start a follow up PR porting our work from ml-gde/jflux

[monatis]

Hi @SauravMaheshkar, thanks for taking a snab to this. It sounds like a plausable idea! And perhaps this modeling from Transformers might be helpful for T5 --you can choose to port it to this project. Feel free to reach out to me any time if you think I might be helpful in any way.

[SauravMaheshkar]

Sounds good. Opened another issue in #34.