TF: XLA stable softmax
Author: gante
Pull request: yes

# What does this PR do?

As discussed in the thread about XLA problems (https://github.com/huggingface/transformers/issues/16838), this PR adds a stable wrapper for the softmax operation, and replaces `tf.nn.softmax` by the wrapped function.

This PR:
- Adds the wrapped softmax, named `stable_softmax`, in `tf_utils.py`. Its docstring includes why it is needed and why the new operation is valid;
- Adds tests to the wrapped softmax, including XLA tests;
- Replaces `tf.nn.softmax` by `stable_softmax` everywhere except in the doctests (I think it overcomplicates the examples, and no XLA should be needed there);
- Removes the `skipIf` for XLA tests, as they can now be successfully executed in a CPU.

Closes #16838 

Comments:
_The documentation is not available anymore as the PR was closed or merged._

This looks good to me! Do you think it would be better to change `stable_softmax` to only add the offset if we're running on CPU? It makes very little difference either way, but we could hide the complexity of that inside `stable_softmax` and keep our code paths entirely unchanged on GPU. I'm not certain, though - since it's such a small change maybe we can just do it everywhere.

> This looks good to me! Do you think it would be better to change `stable_softmax` to only add the offset if we're running on CPU? It makes very little difference either way, but we could hide the complexity of that inside `stable_softmax` and keep our code paths entirely unchanged on GPU. I'm not certain, though - since it's such a small change maybe we can just do it everywhere.

Good point! Hope this won't affect tests on GPU (at least not for PT/TF equivalence which use `1e-5`). Let's see!

@Rocketknight1 @ydshieh if you run the test and print the difference between `stable_softmax` and `tf.nn.softmax`, the difference is exactly `0.0` -- I don't think we need to worry about that :D

@gante With this, do we still have issues regarding sampling in `generate()`. Sorry, I didn't really follow that issue about sampling, but would like to know a bit more 😄 

@ydshieh after this fix, the errors related to `generate()` are gone -- they were caused by the forward pass in the models, which in turn were caused by the issue this PR solves

(I might be completely wrong below)

I could imagine that we (will) have tests like:

- testing non-XLA and XLA `generte()` that use sampling
  - even with this PR, the differences of output logits between these two might still be as large as, say, `1e-3`?
  - if so, the sampling might give different sampling results ..?
  - if not, what's the magnitude of the diff we get after this PR?
 - testing PT and TF `generte()` that use sampling
   - so same potential issue as above ..? 

Thanks 🙏 

OK, I saw your previous comment

```
I've spun up an Nvidia T4 ( = no tf32 format) and got an error < 1e-5 for all cases
```

Based on the testing results, I'm happy for this to be merged now! If this is an XLA bug, though, we should make sure to revert our changes once none of the TF versions we support are affected by it anymore.

Should we add a TODO to the `masked_softmax` function or a reminder somewhere to make sure that we document why this change is here, and when it can be removed?

@Rocketknight1 added a TODO with instructions related to when to deprecate 👍