tf.function and half precision fails with Roberta models
Author: AWilcke
Pull request: no

## Environment info

- `transformers` version: 4.5.1
- Platform: Linux-5.4.0-71-generic-x86_64-with-glibc2.29
- Python version: 3.8.5
- PyTorch version (GPU?): 1.7.1 (False)
- Tensorflow version (GPU?): 2.4.1 (False)
- Using GPU in script?: No (but also fails with GPU)
- Using distributed or parallel set-up in script?: No

### Who can help
As far as I can tell, this worked before #9788, so maybe @jplu can help. Also this is a TF issue so @Rocketknight1 .

## Information

Model I am using (Bert, XLNet ...): TFRoberta, this also happens with TFXLMRoberta

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

```py3
import tensorflow as tf
from transformers.models.roberta import RobertaTokenizerFast, TFRobertaModel


@tf.function
def get_embeddings(
    tokenizer: RobertaTokenizerFast, model: TFRobertaModel, text: str
) -> tf.Tensor:
    return model(**tokenizer(text, return_tensors="tf")).last_hidden_state


if __name__ == "__main__":
    tf.keras.mixed_precision.set_global_policy("float16")
    name = "roberta-base"
    tokenizer = RobertaTokenizerFast.from_pretrained(name)
    model = TFRobertaModel.from_pretrained(name)
    embeddings = get_embeddings(
        tokenizer=tokenizer,
        model=model,
        text="tf.function and mixed precision",
    )
    print(embeddings)
```

Traceback:
```
  File "roberta_bug.py", line 17, in <module>
    embeddings = get_embeddings(
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 828, in __call__
    result = self._call(*args, **kwds)
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 871, in _call
    self._initialize(args, kwds, add_initializers_to=initializers)
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 725, in _initialize
    self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 2969, in _get_concrete_function_internal_garbage_collected
    graph_function, _ = self._maybe_define_function(args, kwargs)
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3361, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3196, in _create_graph_function
    func_graph_module.func_graph_from_py_func(
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 990, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 634, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 977, in wrapper
    raise e.ag_error_metadata.to_exception(e)
TypeError: in user code:

    roberta_bug.py:9 get_embeddings  *
        return model(**tokenizer(text, return_tensors="tf")).last_hidden_state
    /home/arthur/reinfer/env/lib/python3.8/site-packages/transformers/models/roberta/modeling_tf_roberta.py:744 call  *
        outputs = self.roberta(
    /home/arthur/reinfer/env/lib/python3.8/site-packages/transformers/models/roberta/modeling_tf_roberta.py:544 call  *
        extended_attention_mask = tf.multiply(tf.subtract(1.0, extended_attention_mask), -10000.0)
    /home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:201 wrapper  **
        return target(*args, **kwargs)
    /home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py:561 subtract
        return gen_math_ops.sub(x, y, name)
    /home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/ops/gen_math_ops.py:10316 sub
        _, _, _op, _outputs = _op_def_library._apply_op_helper(
    /home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/framework/op_def_library.py:555 _apply_op_helper
        raise TypeError(

    TypeError: Input 'y' of 'Sub' Op has type float16 that does not match type float32 of argument 'x'.
```

## Expected behavior

The model should calculate embeddings correctly. This is due to `tf.subtract(1.0, extended_attention_mask)` checking that `1.0` and `extended_attention_mask` have the same type, but in `float16` mode they do not. Reverting to `1.0 - extended_attention_mask` fixes the issue.


Comments:
Thanks for sharing this issue!

The issue here indeed comes from 1.0 that is not from the same dtype than `extended_attention_mask`. A better fix would be to align this line such as the other models by replacing it with:
```
extended_attention_mask = tf.cast(extended_attention_mask, dtype=embedding_output.dtype)
one_cst = tf.constant(1.0, dtype=embedding_output.dtype)
ten_thousand_cst = tf.constant(-10000.0, dtype=embedding_output.dtype)
extended_attention_mask = tf.multiply(tf.subtract(one_cst, extended_attention_mask), ten_thousand_cst)
```
Extracted from BERT. I will do a fix.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.