Conditional branching logic in modeling_tf_xlnet.py causing error with TF Graph
Author: ANarayan
Pull request: no
Labels: name: wontfix

Hi  @TevenLeScao ,
 
I am encountering an error when running the TFXLNet model inside of a tensorflow graph.

Here is some code to reproduce the issue:

```
from transformers import XLNetTokenizer, TFXLNetModel
import tensorflow as tf
tokenizer = XLNetTokenizer.from_pretrained('xlnet-base-cased')
model = TFXLNetModel.from_pretrained('xlnet-base-cased')

@tf.function
def train_step(inputs, mask, token_type_ids):
    with tf.GradientTape() as tape:
        a = model({
            "input_ids": inputs,
            "training": True,
            "attention_mask": mask,
            "token_type_ids": token_type_ids,
        })
        
train_step(inputs, mask, token_type_ids)
```

The error seems to be caused by L765-L768 in modeling_tf_xlnet.py [here](https://github.com/huggingface/transformers/blob/82498cbc37d5c15520c7bddde5d804c804eee498/src/transformers/models/xlnet/modeling_tf_xlnet.py#L765)

Here is the error message:

>  TypeError: in user code:
>    <ipython-input-41-b79f96ef9347>:4 train_step  *
>        a = model({
>     /anaconda2/envs/dialogue/lib/python3.7/site-packages/transformers/models/xlnet/modeling_tf_xlnet.py:1189 call  *
>         outputs = self.transformer(
>     /anaconda2/envs/dialogue/lib/python3.7/site-packages/transformers/models/xlnet/modeling_tf_xlnet.py:753 call  *
>         if inputs["use_mems"]:
>     /anaconda2/envs/dialogue/lib/python3.7/site-packages/tensorflow/python/autograph/operators/control_flow.py:951 if_stmt
>         _tf_if_stmt(cond, body, orelse, get_state, set_state, symbol_names, nouts)
>     /anaconda2/envs/dialogue/lib/python3.7/site-packages/tensorflow/python/autograph/operators/control_flow.py:996 _tf_if_stmt
>         cond, aug_body, aug_orelse, strict=True)
>     /anaconda2/envs/dialogue/lib/python3.7/site-packages/tensorflow/python/util/dispatch.py:201 wrapper
>         return target(*args, **kwargs)
>     /anaconda2/envs/dialogue/lib/python3.7/site-packages/tensorflow/python/util/deprecation.py:507 new_func
>         return func(*args, **kwargs)
>     /anaconda2/envs/dialogue/lib/python3.7/site-packages/tensorflow/python/ops/control_flow_ops.py:1180 cond
>         return cond_v2.cond_v2(pred, true_fn, false_fn, name)
>     /anaconda2/envs/dialogue/lib/python3.7/site-packages/tensorflow/python/ops/cond_v2.py:92 cond_v2
>         op_return_value=pred)
>     /anaconda2/envs/dialogue/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py:986 func_graph_from_py_func
>         func_outputs = python_func(*func_args, **func_kwargs)
>     /anaconda2/envs/dialogue/lib/python3.7/site-packages/tensorflow/python/autograph/operators/control_flow.py:992 aug_orelse
>         _verify_tf_cond_vars(new_body_vars_[0], new_orelse_vars, symbol_names)
>     /anaconda2/envs/dialogue/lib/python3.7/site-packages/tensorflow/python/autograph/operators/control_flow.py:286 _verify_tf_cond_vars
>         ' branches:\n\n{}'.format(name, str(e)))
>     TypeError: 'new_mems' must have the same nested structure in the main and else branches:
>     The two structures don't have the same nested structure.
>     First structure: type=tuple str=(<tf.Tensor 'tfxl_net_model/transformer/cond_2/StopGradient:0' shape=(44, 18, 768) dtype=float32>,)
>     Second structure: type=tuple str=()
>     More specifically: The two structures don't have the same number of elements. First structure: type=tuple str=(<tf.Tensor 'tfxl_net_model/transformer/cond_2/StopGradient:0' shape=(44, 18, 768) dtype=float32>,). Second structure: type=tuple str=()
>     Entire first structure:
>     (.,)
>     Entire second structure:
>     ()

Comments:
This issue has been automatically marked as stale and been closed because it has not had recent activity. Thank you for your contributions.

If you think this still needs to be addressed please comment on this thread.