implementing flux on TPUs with ptxla

[entrpn]

What does this PR do?

Flux inference using PyTorch XLA.

I still need to validate that the readme works on a new machine and I didn't miss any steps then run the linter. Will do it tomorrow.

@pcuenca @sayakpaul I have a question regarding the AttentionProcessors. In a previous PR it was recommended to create a new AttentionProcessor for ptxla. Unfortunately that one doesn't work for Flux and I feel its a lot of duplication of code if I do the same for Flux as it would almost be a copy/paste of FluxAttentionProcessors. Please advice.

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[sayakpaul]

it was recommended to create a new AttentionProcessor for ptxla. Unfortunately that one doesn't work for Flux and I feel its a lot of duplication of code if I do the same for Flux as it would almost be a copy/paste of FluxAttentionProcessors. Please advice.

Okay for me to create a new class even if that means duplication. Once we start to more formally support PyTorch XLA (i.e., through inclusion in our CI), we can start revisiting them to refactor but for now, new classes are completely fine by me. Cc: @yiyixuxu WDYT?

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[entrpn]

@hlky @sayakpaul created new class and addressed the comments + ran make style/quality. Please take another look.

[sayakpaul]

@entrpn could you run make style && make quality?

[entrpn]

make style && make quality

done

[yiyixuxu]

thanks!

[sayakpaul]

Thanks for contributing this!