Support beam search

Hi,

It would be nice to support beam search.

There is [the reference flax implementation in wmt example](https://github.com/google/flax/blob/main/examples/wmt/decode.py) and [the equivalent one from `transformers`](https://github.com/huggingface/transformers/blob/main/src/transformers/generation/flax_utils.py).

I am guessing that we could:
* duplicate inputs per num_beams initially
* at each step we do:
* decode_step
* select top beams
* overwrite entire past cache per selected beams
* update cache with new selected tokens

So maybe the extra step here is to add the "overwrite entire past cache per selected beams"?
Curious if you have suggestions for implementation