Return scalar losses instead of per-sample means
This updates the TF XLA-compatible losses to return scalars instead of per-sample means. As @ydshieh pointed out, per-sample means give too much weight to samples with fewer masked positions. The new approach should match PyTorch losses exactly (up to floating-point error).

TODO:
- [x] Update expected sizes in tests
_The documentation is not available anymore as the PR was closed or merged._<|||||>@ydshieh will investigate RAG and make sure tests pass before merging!<|||||>@ydshieh I reverted the RAG loss function to the pre-XLA version, so hopefully those tests pass now. All other tests are passing! Do you think there's anything else I'm missing before I merge?<|||||>@gante We do, but it [takes the mean of the TF loss in order to make the two losses comparable](https://github.com/huggingface/transformers/blob/main/tests/test_modeling_tf_common.py#L560-L567).<|||||>> Do you think there's anything else I'm missing before I merge?

Nothing I can think of as the tests pass now 🎉 . Thank you!