TF: all models can run in Graph mode
Author: gante
Pull request: no
Labels: name: WIP

### Feature request

Our models are executed in Eager mode by default. Eager mode is more permissive than Graph mode, and we have models that don't work in Graph mode at the moment. This (self) feature request is being added to bring visibility to the problem, link issues, and track progress.

### Motivation

It is a requirement for several downstream uses, like XLA-accelerated forward passes or TF serving.

### Your contribution

Adding a general test to ensure all existing and new models are compatible with graph mode. I will also attempt to fix as many related issues as I can.

Comments:
cc @Rocketknight1 and @ydshieh -- link related issues as you see them plz :D 

Strong +1 on this, since compiling in graph mode is needed for all the other downstream things we want to do (TF Serving, tf.js, etc.)

Link a potentially related issue as early as possible, as I have a poor memory capability

https://github.com/huggingface/transformers/pull/16886#issuecomment-1113448810

[#17233](https://github.com/huggingface/transformers/issues/17233)

> [#17233](https://github.com/huggingface/transformers/issues/17233)

To add another info: I found we have something like the following in `TFWav2Vec2Encoder`
```
# add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
dropout_probability = np.random.uniform(0, 1)
if training and (dropout_probability < self.config.layerdrop):  # skip the layer
    continue
```
I think we should avoid this, right?

@ydshieh it's undesirable, but I think TF can handle it (with potential retracing if the parameters change). Will have to double-check, though

I do have a general fear with things like `LayerDrop`, which is that XLA cannot accept any data-dependent computation paths. In other words, you cannot have a scenario where a layer is only run if a random number, generated by the GPU for each sample/batch, is over a threshold value. You can only implement something like this by running the layer every time with a residual connection, and multiplying the layer outputs by 0 if it is going to be "dropped". Doing this, of course, has no performance benefit at all.

In code like the above, the number is generated by `numpy`, which will be run once at the point of graph tracing. Therefore, that layer will either be skipped *always* or *never*. Correct code would use `tf.random` instead, which will insert the random generation into the graph correctly, but then the `LayerDrop` would cause XLA tracing to fail, and I'm not sure about regular Graph mode tracing.

> you cannot have a scenario where a layer is only run if a random number

`tf.cond` should be able to handle it, I believe 🤔  As you mentioned, with further changes as well. This is going to be fun! 

Bonus: the just-released TF 2.9 compiles the forward pass with XLA when `Model.compile(jit_compile=True)` is called, so we might be able to get cool performance numbers out of sorting this on as many models as we can 😎 

Ah, I'm completely wrong - I was confusing two of the XLA requirements. `tf.cond()` will solve it all, I'm sorry!

Hey has this been fixed? 

Hey @ahmedlone127 👋 It hasn't been fixed yet.

Okay thanks :) 

#18153 Ensures all our models can be saved as `SavedModel`, with the exception of CLIP. I'm closing this issue as having a `SavedModel` implies this issue is solved 

Kudos to @amyeroberts for smashing it!