Fix loss computation in TFBertForPreTraining
With thanks to @sreyan88 for writing up a clean bug report and reproducer, and to @ydshieh for locating the problematic code!

Our `hf_compute_loss()` function for `TFBertForPreTraining` was incorrect. However, it still appeared to work when the number of masked positions was evenly divisible by the batch size. Other, more commonly-used models like `TFBertForMaskedLM` do not have this issue.

The problem was incorrect handling of the reduction for the masked loss, so I took the opportunity to rewrite the function in modern TF. All shapes are now static in the rewritten function as well, which means it should now compile with XLA.

Fixes #17883 
_The documentation is not available anymore as the PR was closed or merged._<|||||>@sgugger I believe it's unique to BERT, because I tried searching the codebase for any similar lines and it couldn't find any. I suspect this is how it stayed undetected for so long - it uses the NSP loss and people generally don't train with that anymore.