Adding XLA support for greedy sampling
# What does this PR do?

This CR enables greedy sampling in model.generate on XLA devices such as Trainium and TPU. This addresses issues such as  https://github.com/huggingface/transformers/issues/18661 and https://github.com/huggingface/transformers/issues/12322. 

The implementation is inspired by the corresponding Tensorflow generate function in transformers. The CR uses conditional statements to support greedy sampling, and the user can switch between GPU implementation or XLA implementation depending on the state of `is_torch_tpu_available`. The CR also implements kv-cache functionality that is XLA compatible.

Fixes # (issue)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ X] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [X ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? 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](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@gante Feel free to suggest appropriate tests/refactors for this PR. We have tested generation locally using a trn1.32xlarge instance and matched Rouge scores for T5-small summarization.

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_23021). All of your documentation changes will be reflected on that endpoint.<|||||>Hey @aashiqmuhamed! In general the PR looks positive, but before diving deeper, let us (`transformers` team) have a discussion about adding this type of PRs (new HW-oriented optimizations). `generate` is very complex ATM, and we want to make it more manageable -- if we merge all PRs of this kind in the current `generate` state, it will become maintenance hell for everyone. 

I will get back to you in this PR within a week.<|||||>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.