Add Flax implementation for BLOOM
### Model description

I'm interested in adding an implementation of BLOOM in Flax. 

The implementation shouldn't be too bad since the pytorch implementation can serve as a guide and a way to check correctness. 



### Open source status

- [X] The model implementation is available
- [X] The model weights are available

### Provide useful links for the implementation

@younesbelkada @stas00 @patrickvonplaten 

If someone is already planning to work on this then no worries, but if not I will start on this as soon as I have time!
Hi!
Thank you very much for the contribution!
On my side it's a green light since I am not working on it, and it is not on my plans for now. Therefore, I'll be happy to review it! Let us know if you want to work on that :)<|||||>Thanks! I will open a WIP PR soon and tag you there once I do.<|||||>Very cool idea - think this can also be a flagship project where we can showcase how to fine-tune BLOOM with Flax cc @patil-suraj @sanchit-gandhi <|||||>Awesome!  Would be very happy to help with it :) <|||||>Great idea! Would also be interested in getting involved, this would be a super cool model addition!<|||||>Thanks everyone for the interest! I'd love to collaborate with you all.

I'm hoping to push a rough draft of modeling code by the end of the weekend (earlier if I have time), and will tag you all when I open the PR with that. Does that sound alright?<|||||>I've opened a PR (and documented the state of the in-progress code I'm still working on) at #17761 ! We can discuss further in that PR how to collaborate / proceed.