[Flax] Fix incomplete batches in example scripts
# What does this PR do?

Currently in our Flax examples scripts, we drop the last incomplete batch during training and inference:
https://github.com/huggingface/transformers/blob/09178705101b9803e7b9ea7f79a46b4c242dd4bf/examples/flax/summarization/run_summarization_flax.py#L350
We do this for two reasons:

1. Because XLA is not shape polymorphic, forming the last batch with shape different from the preceding batches triggers a recompilation of the `pmap`'d function .
2. If the batch size is not divisible by the number devices, then the last step must be executed on a single device (or a subset of devices), potentially leading to OOMs.

During training, dropping the last batch isn't an issue: since we shuffle the data and train for multiple epochs, all of the training data is eventually used and the effects of dropping the last batch amortised.

However, during evaluation and prediction, dropping the last batch leads to incorrect results: since we don't account for the examples in the last batch, we do not evaluate over the whole dataset, and thus have partial results.

This PR corrects for the incomplete batches in the relevant Flax training examples.

## 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?
- [ ] 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?

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.
_The documentation is not available anymore as the PR was closed or merged._<|||||>Just waiting to double check that the slow tests pass from [`test_flax_examples.py`](https://github.com/huggingface/transformers/blob/main/examples/flax/test_flax_examples.py) before merging. Working with @patil-suraj to verify this ✅<|||||>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.<|||||>@patil-suraj @sanchit-gandhi can we merge this one?<|||||>Just verifying the slow tests from [`test_flax_examples.py`](https://github.com/huggingface/transformers/blob/main/examples/flax/test_flax_examples.py) pass on a <s>v3-8</s> v100 GPU!