🐛 [TFBART] LayerDrop not working on TPU
Author: astariul
Pull request: no

## 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 TFBart on TPU, I'm getting the following error :

> ValueError: "attn" is None at the end of the TRUE branch.

It seems to come from the LayerDrop operation :

https://github.com/huggingface/transformers/blob/b01ddc9577b87f057e163d49563ee3f74f4810cf/src/transformers/models/bart/modeling_tf_bart.py#L387-L391

<details>
 <summary> Full stack trace (click to expand...)</summary>

>2020/12/11 00:35:34 - INFO - transformers_addons.trainer_tf -   ***** Running Evaluation *****
2020/12/11 00:35:34 - 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)
ValueError: 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:388 call  *
        if training and (dropout_probability < self.layerdrop):  # skip the layer
    /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:242 _verify_single_cond_var
        raise ValueError('"{}" is None at the end of the TRUE branch.'.format(name))
>
>    ValueError: "attn" is None at the end of the TRUE branch.

</details>


Comments:
It seems to work if I completely remove the `LayerDrop` (by commenting out the `if` clause, in both encoder and decoder).

Hey @astariul-colanim, 

I think a fix in #9029 (replacing the if-else by a `continue` statement) should do the trick.

Could you try again from the branch and let me know?

Thanks!

So far #9029 seems working perfectly !
Let's close this issue when #9029  is merged :)

Thanks for the fix !

@patrickvonplaten Finally the model crash during evaluation..

<details>
<summary> Full stack trace (click to view)</summary>

```
2020/12/15 01:19:09 - INFO - transformers_addons.trainer_tf -   ***** Running Evaluation *****
2020/12/15 01:19:09 - 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=Tru
e)`).
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=Tru
e)`).
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=Tru
e)`).
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 190, in main
    trainer.train()
  File "/home/remondnicola/text-summarization/transformers_addons/trainer_tf.py", line 339, in train
    result = self.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)
ValueError: 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:97 call  *
	outputs = super().call(inputs["input_ids"],
/home/remondnicola/.venv/summarization/lib/python3.7/site-packages/transformers/models/bart/modeling_tf_bart.py:1222 call  *
	outputs = self.model(
/home/remondnicola/.venv/summarization/lib/python3.7/site-packages/transformers/models/bart/modeling_tf_bart.py:1062 call  *
	inputs["encoder_outputs"] = self.encoder(
/home/remondnicola/.venv/summarization/lib/python3.7/site-packages/transformers/models/bart/modeling_tf_bart.py:719 call  *
	for encoder_layer in self.layers:
/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:242 _verify_single_cond_var
	raise ValueError('"{}" is None at the end of the TRUE branch.'.format(name))   

ValueError: "all_attentions" is None at the end of the TRUE branch.
```

</details>

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.