top-k sampling for Flax models
# 🚀 Feature request

We're using a custom (`Flax`) Seq2Seq model (`Bart`) in our `DALL·E mini` project to generate the image tokens, currently, we're doing the following to generate samples (`encoded image`) given a tokenized prompt to the Seq2Seq model:
```
model.generate(
  **tokenized_prompt, 
  do_sample=True, 
  num_beams=1, 
  prng_key=subkey
)
```

But, now we're trying to experiment with the generation method (for example `top_k` sampling) to see if it improves our generated samples. 

We've tried doing the following:
```
model.generate(
        **tokenized_prompt,
        do_sample=True,
        top_k=50,
        prng_key=key,
        params=params
    )
```
Which throws- `NotImplementedError: Beam sampling is currently not implemented.`
After looking at the source-code for `generation_flax_utils`, it's indeed the case that it hasn't been implemented yet ([here](https://huggingface.co/transformers/_modules/transformers/generation_flax_utils.html)).

I hope the sampling feature will be integrated into the flax models, our project and other (Flax ones) will benefit greatly from the feature (and I'm sure it's already in the todo-list of the amazing HF dev team).
[_project repo_](https://github.com/borisdayma/dalle-mini)

<!-- A clear and concise description of the feature proposal.
     Please provide a link to the paper and code in case they exist. -->

## Motivation

We believe that the `top_k` sampling will improve our generations (quite a bit if not a lot), and even in the taming transformers paper, they've used `top_k` sampling for generating tokens with their autoregressive models and have shown higher `top_k` (subjective to the dataset) improves image generations by increasing variance (in a sense). So we just want to experiment with this sampling method and see if we can find something similar for our model.

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/master/CONTRIBUTING.md) are likely to be ignored.