Add Flax diverse group search

### Feature request Add diverse beam search decoding to Flax, an "alternative to BS that decodes a list of diverse outputs by optimising for a diversity-augmented objective", as described in the paper: https://arxiv.org/pdf/1610.02424.pdf This feature would mimic the PyTorch equivalent, added in #9006. @yeandy made a great start on adding this feature in the PR #24508. The PR is still open, and anyone in the community is free to pick-up the PR and see it through to completion! ### Motivation There's a promising PR for this feature that is partway there - it would be a shame not to see this through to completion! ### Your contribution Happy to answer any questions/queries on the PR and provide PR reviews 🤗 Think this would be a fun one for any Flax contributors who are interested!