Repository: huggingface/transformers

Issue:
🐛 [TF_BART] "<internal expr>" has dtype float32 in the TRUE branch, but dtype=int32 in the FALSE branch
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: 4.1.0.dev0
- Platform: Linux-4.19.0-13-cloud-amd64-x86_64-with-debian-10.7
- Python version: 3.7.3
- PyTorch version (GPU?): 1.7.0 (False)
- Tensorflow version (GPU?): 2.2.0 (False)
- Using GPU in script?: No (TPU)
- Using distributed or parallel set-up in script?: Yes

### Who can help
 
@patrickvonplaten

## Information

Model I am using (Bert, XLNet ...): TFBart

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

## To reproduce

When I try to run TF_Bart on TPU, I'm getting the following error :

> TypeError: "<internal expr>" has dtype float32 in the TRUE branch, but dtype=int32 in the FALSE branch. TensorFlow control flow requires that they are the same.

It seems to come from the dropout operation :

https://github.com/huggingface/transformers/blob/b01ddc9577b87f057e163d49563ee3f74f4810cf/src/transformers/models/bart/modeling_tf_bart.py#L373

<details>
 <summary> Full stack trace (click to expand...)</summary>

>2020/12/11 00:00:55 - INFO - transformers_addons.trainer_tf -   ***** Running Evaluation *****
2020/12/11 00:00:55 - INFO - transformers_addons.trainer_tf -     Batch size = 8
The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).
The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.
The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).
The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.
Traceback (most recent call last):
  File "train.py", line 203, in <module>
    main()
  File "train.py", line 194, in main
    result = trainer.evaluate()
  File "/home/remondnicola/text-summarization/transformers_addons/trainer_tf.py", line 281, in evaluate
    output = self._prediction_loop(eval_dataset, description="Evaluation", prediction_loss_only=prediction_loss_only)
  File "/home/remondnicola/text-summarization/transformers_addons/trainer_tf.py", line 207, in _prediction_loop
    loss, logits = self._evaluate_steps(features, labels)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 580, in __call__
    result = self._call(*args, **kwds)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 627, in _call
    self._initialize(args, kwds, add_initializers_to=initializers)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 506, in _initialize
    *args, **kwds))
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py", line 2446, in _get_concrete_function_internal_garbage_collected
    graph_function, _, _ = self._maybe_define_function(args, kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py", line 2777, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py", line 2667, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py", line 981, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 441, in wrapped_fn
    return weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py", line 3299, in bound_method_wrapper
    return wrapped_fn(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py", line 968, in wrapper
    raise e.ag_error_metadata.to_exception(e)
TypeError: in user code:
>
>    /home/remondnicola/text-summarization/transformers_addons/trainer_tf.py:169 _evaluate_steps  *
        per_replica_loss, per_replica_logits = self.args.strategy.experimental_run_v2(
    train.py:29 _run_model  *
        out = self.model(features, training=training, **labels)
    /home/remondnicola/text-summarization/transformers_addons/models/bart/modeling_tf_bart.py:88 call  *
        outputs = super().call(inputs["input_ids"],
    /home/remondnicola/.venv/summarization/lib/python3.7/site-packages/transformers/models/bart/modeling_tf_bart.py:1110 call  *
        outputs = self.model(
    /home/remondnicola/.venv/summarization/lib/python3.7/site-packages/transformers/models/bart/modeling_tf_bart.py:977 call  *
        inputs["encoder_outputs"] = self.encoder(
    /home/remondnicola/.venv/summarization/lib/python3.7/site-packages/transformers/models/bart/modeling_tf_bart.py:373 call  *
        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/operators/control_flow.py:924 if_stmt
        basic_symbol_names, composite_symbol_names)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/operators/control_flow.py:962 tf_if_stmt
        error_checking_orelse)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/util/deprecation.py:507 new_func
        return func(*args, **kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/control_flow_ops.py:1177 cond
        return cond_v2.cond_v2(pred, true_fn, false_fn, name)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/cond_v2.py:91 cond_v2
        op_return_value=pred)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py:981 func_graph_from_py_func
        func_outputs = python_func(*func_args, **func_kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/operators/control_flow.py:958 error_checking_orelse
        basic_symbol_names + composite_symbol_names)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/operators/control_flow.py:298 _verify_tf_cond_vars
        functools.partial(_verify_single_cond_var, name), body_var, orelse_var)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/util/nest.py:617 map_structure
        structure[0], [func(*x) for x in entries],
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/util/nest.py:617 <listcomp>
        structure[0], [func(*x) for x in entries],
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/operators/control_flow.py:267 _verify_single_cond_var
        orelse_var.dtype.name))
>
>    TypeError: "<internal expr>" has dtype float32 in the TRUE branch, but dtype=int32 in the FALSE branch. TensorFlow control flow requires that they are the same.

</details>


Discussion:

Resolving patch:
diff --git a/src/transformers/models/bart/modeling_tf_bart.py b/src/transformers/models/bart/modeling_tf_bart.py
--- a/src/transformers/models/bart/modeling_tf_bart.py
+++ b/src/transformers/models/bart/modeling_tf_bart.py
@@ -235,9 +235,9 @@ def __init__(self, config: BartConfig, **kwargs):
         )
         self.normalize_before = config.normalize_before
         self.self_attn_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="self_attn_layer_norm")
-        self.dropout = config.dropout
+        self.dropout = tf.keras.layers.Dropout(config.dropout)
         self.activation_fn = ACT2FN[config.activation_function]
-        self.activation_dropout = config.activation_dropout
+        self.activation_dropout = tf.keras.layers.Dropout(config.activation_dropout)
         self.fc1 = tf.keras.layers.Dense(config.encoder_ffn_dim, name="fc1")
         self.fc2 = tf.keras.layers.Dense(self.embed_dim, name="fc2")
         self.final_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="final_layer_norm")
@@ -261,7 +261,7 @@ def call(self, x, encoder_padding_mask, training=False):
         assert shape_list(x) == shape_list(
             residual
         ), f"Self attn modified the shape of query {shape_list(residual)} to {shape_list(x)}"
-        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
+        x = self.dropout(x, training=training)
         x = residual + x
         if not self.normalize_before:
             x = self.self_attn_layer_norm(x)
@@ -270,9 +270,9 @@ def call(self, x, encoder_padding_mask, training=False):
         if self.normalize_before:
             x = self.final_layer_norm(x)
         x = self.activation_fn(self.fc1(x))
-        x = tf.nn.dropout(x, rate=self.activation_dropout if training else 0)
+        x = self.activation_dropout(x, training=training)
         x = self.fc2(x)
-        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
+        x = self.dropout(x, training=training)
         x = residual + x
         if not self.normalize_before:
             x = self.final_layer_norm(x)
@@ -293,7 +293,7 @@ class TFBartEncoder(tf.keras.layers.Layer):
     def __init__(self, config: BartConfig, embed_tokens: TFSharedEmbeddings, **kwargs):
         super().__init__(**kwargs)
 
-        self.dropout = config.dropout
+        self.dropout = tf.keras.layers.Dropout(config.dropout)
         self.layerdrop = config.encoder_layerdrop
         self.output_hidden_states = config.output_hidden_states
         self.output_attentions = config.output_attentions
@@ -370,7 +370,7 @@ def call(
         embed_pos = self.embed_positions(input_ids)
         x = inputs_embeds + embed_pos
         x = self.layernorm_embedding(x)
-        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
+        x = self.dropout(x, training=training)
 
         # B x T x C -> T x B x C
         x = tf.transpose(x, perm=[1, 0, 2])
@@ -413,9 +413,9 @@ def __init__(self, config: BartConfig, **kwargs):
             dropout=config.attention_dropout,
             name="self_attn",
         )
-        self.dropout = config.dropout
+        self.dropout = tf.keras.layers.Dropout(config.dropout)
         self.activation_fn = ACT2FN[config.activation_function]
-        self.activation_dropout = config.activation_dropout
+        self.activation_dropout = tf.keras.layers.Dropout(config.activation_dropout)
         self.normalize_before = config.normalize_before
 
         self.self_attn_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="self_attn_layer_norm")
@@ -467,7 +467,7 @@ def call(
             attn_mask=causal_mask,
             key_padding_mask=decoder_padding_mask,
         )
-        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
+        x = self.dropout(x, training=training)
         x = residual + x
         if not self.normalize_before:
             x = self.self_attn_layer_norm(x)
@@ -481,7 +481,7 @@ def call(
             key_padding_mask=encoder_attn_mask,
             layer_state=layer_state,  # mutates layer state
         )
-        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
+        x = self.dropout(x, training=training)
         x = residual + x
         if not self.normalize_before:
             x = self.encoder_attn_layer_norm(x)
@@ -490,9 +490,9 @@ def call(
         if self.normalize_before:
             x = self.final_layer_norm(x)
         x = self.activation_fn(self.fc1(x))
-        x = tf.nn.dropout(x, rate=self.activation_dropout if training else 0)
+        x = self.activation_dropout(x, training=training)
         x = self.fc2(x)
-        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
+        x = self.dropout(x, training=training)
         x = residual + x
         if not self.normalize_before:
             x = self.final_layer_norm(x)
@@ -545,7 +545,7 @@ def __init__(self, config: BartConfig, embed_tokens, **kwargs):
             else None
         )
 
-        self.dropout = config.dropout
+        self.dropout = tf.keras.layers.Dropout(config.dropout)
         self.output_hidden_states = config.output_hidden_states
         self.output_attentions = config.output_attentions
         self.use_cache = config.use_cache
@@ -588,7 +588,7 @@ def call(
             x = self.layernorm_embedding(x) + positions
         else:
             x = self.layernorm_embedding(x + positions)
-        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
+        x = self.dropout(x, training=training)
 
         # Convert to Bart output format: (BS, seq_len, model_dim) ->  (seq_len, BS, model_dim)
         x = tf.transpose(x, perm=(1, 0, 2))
@@ -674,7 +674,7 @@ def __init__(
         self.embed_dim = embed_dim
 
         self.num_heads = num_heads
-        self.dropout = dropout
+        self.dropout = tf.keras.layers.Dropout(dropout)
         self.head_dim = embed_dim // num_heads
         assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
         self.scaling = self.head_dim ** -0.5
@@ -772,7 +772,7 @@ def call(
             attn_weights = tf.reshape(attn_weights, (bsz * self.num_heads, tgt_len, src_len))
 
         attn_weights = tf.nn.softmax(attn_weights, axis=-1)
-        attn_probs = tf.nn.dropout(attn_weights, rate=self.dropout if training else 0.0)
+        attn_probs = self.dropout(attn_weights, training=training)
 
         attn_output = tf.matmul(attn_probs, v)  # shape: (bsz * self.num_heads, tgt_len, self.head_dim)
         attn_output = tf.transpose(attn_output, perm=(1, 0, 2))


Test patch: