

TITLE
Conditional branching logic in modeling_tf_xlnet.py causing error with TF Graph

COMMENTS
1

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
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:
>     ()