BLOOM Flax
# What does this PR do?

An attempt of adding Flax implementation of BLOOM - original PR from @haileyschoelkopf #17761 

## TODOs:

- [x] alibi shifting for batched generation
- [ ] change mask fill value cc @patil-suraj 
- [x] optimize code (alibi creation + mask creation)

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_18022). All of your documentation changes will be reflected on that endpoint.<|||||>Is this ready to merge @younesbelkada ? @sanchit-gandhi @patil-suraj could you take a final look?<|||||>Yeah ready to merge for me, for partitioning compatibility I think that we can address another PR for that<|||||>Also we need to add the credits to @haileyschoelkopf for starting the contribution<|||||>Modelling file and tests LGTM. Am in agreement that the the named axes should come as a follow-up PR - better to modify a script that's already in-place and working rather than merging a complete overhaul of how we structure our code in one go.<|||||>Good to merge for me once @patil-suraj and @sanchit-gandhi give the green light here :-) <|||||>Thanks all for your comments! 
I had to fix a small test since the main modeling code has changed a bit in #18344 !
I think this PR is ready to be merged 💪 cc @patil-suraj <|||||>@patil-suraj @sanchit-gandhi - should we merge this one? <|||||>Good for me! Awaiting the green light from @patil-suraj who requested a look over the modelling code prior to merging. Regarding `setup`, I previously asked in the G Chat with the JAX/Flax guys how to make this work with `scan`. Think it went under the radar, so will bring this issue back up!<|||||>Hi @sanchit-gandhi @patil-suraj @patrickvonplaten 
I think that the PR is ready for a final round of review! 💪  I proposed a potential solution to re-write the model definition with `setup` but I am unsure how this fits the HF+Jax approach regarding `scan` feature<|||||>@sanchit-gandhi @patil-suraj could you check here? I won't have time do review the PR this week I think<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.<|||||>Should we merge this one ? cc @patrickvonplaten @patil-suraj @sanchit-gandhi <|||||>The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_18022). All of your documentation changes will be reflected on that endpoint.<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.<|||||>Both https://github.com/huggingface/transformers/pull/17761 and this PR show a lot of work. Why isn't it merged? <|||||>Adding this to my TODOs