New TF output proposal
Author: jplu
Pull request: yes
Labels: name: wontfix

Hello!

Currently the TensorFlow have several issues to properly working in graph mode when `output_attentions`, `outout_hidden_states` and `return_dict` takes a boolean tensor (`tf.constant(True/False)`). This is because the graph mode doesn't allows to have an output of different sizes in its conditional branches.

To fix this and lower as much as possible a breaking change with the current behavior of these models, I propose to keep the current behavior but only in eager mode, and disabling these three features (force to have them to False) in graph mode. This is for me the best compromise between having something that works in graph mode and not having a big breaking change.

The graph mode is most of the time used to get fast training/inference and not for doing experiments, and I don't see the point to deploy a model in production that gives you all the attentions/hidden states values, or possibly getting an OOM during the training/inference mostly for the new TFLongformer model.

Here some examples.

Example 1: run in eager mode with default parameters value
```python
from transformers import TFBertForMaskedLM

model = TFBertForMaskedLM.from_pretrained("bert-base-cased")
model(tf.constant([[10,11,12]]))

# Result
(<tf.Tensor: shape=(1, 3, 28996), dtype=float32, numpy=
array([[[ -9.900146 , -10.345711 ,  -9.721893 , ...,  -7.7057724,
          -8.035212 ,  -8.271875 ],
        [ -7.866758 ,  -8.147887 ,  -7.864182 , ...,  -6.7170696,
          -6.2896423,  -6.79779  ],
        [ -9.790001 , -10.2271185,  -9.584716 , ...,  -7.539982 ,
          -7.807548 ,  -8.135937 ]]], dtype=float32)>,)
```

Example 2: run in eager mode by updating a boolean parameter value
```python
from transformers import TFBertForMaskedLM

model = TFBertForMaskedLM.from_pretrained("bert-base-cased")
model(tf.constant([[10,11,12]]), return_dict=True)

# Result
TFMaskedLMOutput(loss=None, logits=<tf.Tensor: shape=(1, 3, 28996), dtype=float32, numpy=
array([[[ -9.900146 , -10.345711 ,  -9.721893 , ...,  -7.7057724,
          -8.035212 ,  -8.271875 ],
        [ -7.866758 ,  -8.147887 ,  -7.864182 , ...,  -6.7170696,
          -6.2896423,  -6.79779  ],
        [ -9.790001 , -10.2271185,  -9.584716 , ...,  -7.539982 ,
          -7.807548 ,  -8.135937 ]]], dtype=float32)>, hidden_states=None, attentions=None)
```

Example 3: run in graph mode with default parameters value
```python
from transformers import TFBertForMaskedLM
import tensorflow as tf

model = tf.function(TFBertForMaskedLM.from_pretrained("bert-base-cased"))
model(tf.constant([[10,11,12]]))

# Result
(<tf.Tensor: shape=(1, 3, 28996), dtype=float32, numpy=
array([[[ -9.900146 , -10.345711 ,  -9.721893 , ...,  -7.7057724,
          -8.035212 ,  -8.271875 ],
        [ -7.866758 ,  -8.147887 ,  -7.864182 , ...,  -6.7170696,
          -6.2896423,  -6.79779  ],
        [ -9.790001 , -10.2271185,  -9.584716 , ...,  -7.539982 ,
          -7.807548 ,  -8.135937 ]]], dtype=float32)>,)
```

Example 4: run in graph mode by updating a boolean parameter value
```python
from transformers import TFBertForMaskedLM
import tensorflow as tf

model = tf.function(TFBertForMaskedLM.from_pretrained("bert-base-cased"))
model(tf.constant([[10,11,12]]), return_dict=tf.constant(True))

# Result
(<tf.Tensor: shape=(1, 3, 28996), dtype=float32, numpy=
array([[[ -9.900146 , -10.345711 ,  -9.721893 , ...,  -7.7057724,
          -8.035212 ,  -8.271875 ],
        [ -7.866758 ,  -8.147887 ,  -7.864182 , ...,  -6.7170696,
          -6.2896423,  -6.79779  ],
        [ -9.790001 , -10.2271185,  -9.584716 , ...,  -7.539982 ,
          -7.807548 ,  -8.135937 ]]], dtype=float32)>,)
```
As you can see for the last example, no more graph compilation error due to the `tf.constant(True)` value.

For my examples I used only `return_dict` but of course the same behavior is applied to `output_attentions` and `output_hidden_states`. Moreover, this new behavior will be detailed in the documentation of the impacted parameters.

@LysandreJik @sgugger @patrickvonplaten @sshleifer @thomwolf @julien-c What do you think of this solution?

Comments:
This proposal works for me. It's better to allow graph mode for the "standard" way to do things while keeping backwards compatibility imo. Users that want to leverage hidden states/attentions with graph mode should have no problems in doing a custom architecture that does exactly that. 

However, as a user I would feel like the last example is slightly misleading; in eager mode providing the input `return_dict` changes things, while in graph mode it doesn't, with no information whatsoever. I would expect the model to at least throw a warning telling me that this isn't supported in graph mode rather than silently doing nothing.

Having a warning message for these cases is a great idea! I take care of that to have this added in the next commit!!

Now everytime `return_dict`, `output_attentions` or `output_hidden_states` will be different of None, the following message will be displayed:

```
Warning: The parameters return_dict, output_attentions and output_hidden_states are disabled in graph mode.
```

@LysandreJik Does it looks ok for you now?

Just to let you know, usual logging is disabled in graph mode, so only way you have to display a message is to use the internal `tf.print()` function.

@sgugger Just pushed a new commit, let me know if I forgot to address one of your comments.

I only commented on the first two models, but the same is applicable for all the other ones. Other than that, it looks good yes, thanks for addressing!

Thanks for taking care of this! I like the approach

@sshleifer @thomwolf @julien-c Do you have any comments?

@sgugger @patrickvonplaten and @LysandreJik looks to be agree on that.

I will continue to work on this by applying it to the other models waiting @sshleifer @thomwolf and @julien-c comments.

Fine with me! In general, no need to wait for my input on tf decisions :)

I prefer to have the opinion of everyone as it is also a bit a design update :)

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.