Help with Using TFXLNet on custom embeddings
# 🐛 Bug

## Information

Hi, I am working on implementing the multimodal bitransformer: https://arxiv.org/pdf/1909.02950.pdf
I have already gotten this working using your implementation of TFBertModel, but I want to try using TFXLNetModel in place of BERT to see if it makes an improvement. What I've done for using TFBertModel is extract the word_embeddings and pass the word embeddings with token_type_ids = 0 and image embeddings with token_type_ids = 1. 

## To reproduce

I have the following definition of my model

`
class BERT(transformers.TFXLNetModel):
    def __init__(self, config, *inputs, **kwargs):
        super(BERT, self).__init__(config, *inputs, **kwargs)
        self.call = tf.function(self.call)


class MyModel(tf.keras.Model):
    
    def __init__(self):
        super(MyModel, self).__init__()
        self.resnet = tf.keras.applications.ResNet152V2(include_top=False, 
                                                             weights='imagenet', 
                                                             input_shape=(224, 224, 3))
        self.bert = BERT.from_pretrained('xlnet-base-cased')
        
        self.text_embedding = self.bert.get_input_embeddings().weights[0]
         
        self.pooling = layers.AveragePooling2D(pool_size=(2, 2), padding='same')
        self.reshape = layers.Reshape((4 * 4, 2048)) # 4 is from 7//2 + 1
        self.W_ns = [layers.Dense(self.bert.config.hidden_size) for _ in range(self.reshape.target_shape[0])]
        
        self.concat = layers.Concatenate(axis=1)
        
        self.dropout = layers.Dropout(0.1)
        self.denseout = layers.Dense(1, activation='sigmoid')

    def call(self, inputs):
        text, image = inputs
        
        # handle image
        image = tf.keras.applications.resnet_v2.preprocess_input(image)
        image_emb = self.resnet(image)
        image_emb = self.pooling(image_emb)
        image_emb = self.reshape(image_emb)
        image_embeds = [self.W_ns[i](image_emb[:, i]) for i in range(self.reshape.target_shape[0])]
        image_emb = tf.keras.backend.stack(image_embeds, axis=1)
        
        # handle text
        text_emb = tf.gather(self.text_embedding, text)
        
        # concat and feed to bert
        concat_emb = self.concat([text_emb, image_emb])
        seg_ids = np.concatenate((np.zeros(max_len, dtype=np.int64), 
                                  np.ones(self.reshape.target_shape[0], dtype=np.int64)))
        print('input shapes to xlnet', concat_emb.shape, seg_ids.shape)
        bert_encodings = self.bert(inputs={'inputs_embeds': concat_emb,
                                            'token_type_ids': seg_ids})[0]
        doc_encoding = tf.squeeze(bert_encodings[:, 0:1, :], axis=1)
        doc_encoding = self.dropout(doc_encoding)
        
        output = self.denseout(doc_encoding)
        
        return output
`

In the line that prints "input shapes to xlnet", I get (None, 116, 768) for the inputs_embeds and (116,) for the token_type_ids, which I expect because I have 100 word embeddings and 16 image embeddings. When I call fit() on this model, it gives the error:


> ValueError: in converted code:
> 
>     <ipython-input-12-b6bccd3c83e0>:50 call  *
>         bert_encodings = self.bert(inputs={'inputs_embeds': concat_emb,
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/base_layer.py:891 __call__
>         outputs = self.call(cast_inputs, *args, **kwargs)
>     /homes/awl27/python36env/lib/python3.6/site-packages/transformers/modeling_tf_xlnet.py:824 call  *
>         outputs = self.transformer(inputs, **kwargs)
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/base_layer.py:842 __call__
>         outputs = call_fn(cast_inputs, *args, **kwargs)
>     /homes/awl27/python36env/lib/python3.6/site-packages/transformers/modeling_tf_xlnet.py:530 call  *
>         token_type_ids = tf.transpose(token_type_ids, perm=(1, 0)) if token_type_ids is not None else None
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/ops/array_ops.py:1780 transpose_v2
>         return transpose(a=a, perm=perm, name=name, conjugate=conjugate)
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/ops/array_ops.py:1870 transpose
>         ret = transpose_fn(a, perm, name=name)
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/ops/gen_array_ops.py:11455 transpose
>         "Transpose", x=x, perm=perm, name=name)
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/framework/op_def_library.py:793 _apply_op_helper
>         op_def=op_def)
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/framework/func_graph.py:548 create_op
>         compute_device)
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py:3429 _create_op_internal
>         op_def=op_def)
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py:1773 __init__
>         control_input_ops)
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py:1613 _create_c_op
>         raise ValueError(str(e))
> 
>     ValueError: Dimension must be 1 but is 2 for 'transformer/transpose_1' (op: 'Transpose') with input shapes: [116], [2].
> 

## Expected behavior

I expected this to work just like TFBertModel did. If I just change the definition in the BERT class to use TFBertModel instead of TFXLNetModel, it works fine.

## Environment info
PyTorch version: 1.5.0
Is debug build: No
CUDA used to build PyTorch: 10.2

OS: Ubuntu 16.04.6 LTS
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.11) 5.4.0 20160609
CMake version: version 3.5.1

Python version: 3.6
Is CUDA available: No
CUDA runtime version: 10.0.130
GPU models and configuration: GPU 0: Tesla K40c
Nvidia driver version: 418.87.00
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.3

Versions of relevant libraries:
[pip3] numpy==1.16.6
[pip3] torch==1.5.0
[pip3] torchtext==0.5.0
[pip3] torchvision==0.6.0
[conda] Could not collect

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
