TFXLNet Incompatible shapes in relative attention
Author: h4ste
Pull request: no
Labels: name: wontfix

## 🐛 Bug

TFXLNet fails to run due to incompatible shapes when computing relative attention.

Model I am using (Bert, XLNet....):

TFXLNet

Language I am using the model on (English, Chinese....):

English

The problem arise when using:
* [ ] the official example scripts: (give details)
* [x] my own modified scripts: (give details)

The tasks I am working on is:
* [x] an official GLUE/SQUaD task: *MRPC*
* [ ] my own task or dataset: (give details)

## To Reproduce

Steps to reproduce the behavior:

1. Fix line 542 in `transformers/modeling_tf_xlnet.py` to address #1692:
```python
            input_mask = 1.0 - tf.cast(attention_mask, dtype=dtype_float)
```
2. Run `run_tf_glue_xlnet.py` script:
```python
import os
import tensorflow as tf
import tensorflow_datasets
from transformers import XLNetForSequenceClassification, TFXLNetForSequenceClassification, glue_convert_examples_to_features, XLNetTokenizer

# script parameters
BATCH_SIZE = 32
EVAL_BATCH_SIZE = BATCH_SIZE * 2
USE_XLA = False
USE_AMP = False

# tf.config.optimizer.set_jit(USE_XLA)
# tf.config.optimizer.set_experimental_options({"auto_mixed_precision": USE_AMP})

# Load tokenizer and model from pretrained model/vocabulary
tokenizer = XLNetTokenizer.from_pretrained('xlnet-base-cased')
model = TFXLNetForSequenceClassification.from_pretrained('xlnet-base-cased')

# Load dataset via TensorFlow Datasets
data, info = tensorflow_datasets.load('glue/mrpc', with_info=True)
train_examples = info.splits['train'].num_examples
valid_examples = info.splits['validation'].num_examples

# Prepare dataset for GLUE as a tf.data.Dataset instance
train_dataset = glue_convert_examples_to_features(
        data['train'],
        tokenizer,
        max_length=512,
        output_mode="classification",
        task='mrpc',
        pad_on_left=True, # pad on the left for xlnet
        pad_token=tokenizer.convert_tokens_to_ids([tokenizer.pad_token])[0],
        pad_token_segment_id=4
        )

valid_dataset = glue_convert_examples_to_features(
        data['validation'],
        tokenizer,
        max_length=512,
        output_mode="classification",
        task='mrpc',
        pad_on_left=True, # pad on the left for xlnet
        pad_token=tokenizer.convert_tokens_to_ids([tokenizer.pad_token])[0],
        pad_token_segment_id=4
        )
train_dataset = train_dataset.shuffle(128).batch(BATCH_SIZE).repeat(-1)
valid_dataset = valid_dataset.batch(EVAL_BATCH_SIZE)

# Prepare training: Compile tf.keras model with optimizer, loss and learning rate schedule
opt = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08)
if USE_AMP:
    # loss scaling is currently required when using mixed precision
    opt = tf.keras.mixed_precision.experimental.LossScaleOptimizer(opt, 'dynamic')
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
metric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy')
model.compile(optimizer=opt, loss=loss, metrics=[metric])

# Train and evaluate using tf.keras.Model.fit()
train_steps = train_examples//BATCH_SIZE
valid_steps = valid_examples//EVAL_BATCH_SIZE

history = model.fit(train_dataset, epochs=2, steps_per_epoch=train_steps,
                    validation_data=valid_dataset, validation_steps=valid_steps)

# Save TF2 model
os.makedirs('./save/', exist_ok=True)
model.save_pretrained('./save/')

# Load the TensorFlow model in PyTorch for inspection
pytorch_model = XLNetForSequenceClassification.from_pretrained('./save/', from_tf=True)

# Quickly test a few predictions - MRPC is a paraphrasing task, let's see if our model learned the task
sentence_0 = 'This research was consistent with his findings.'
sentence_1 = 'His findings were compatible with this research.'
sentence_2 = 'His findings were not compatible with this research.'
inputs_1 = tokenizer.encode_plus(sentence_0, sentence_1, add_special_tokens=True, return_tensors='pt')
inputs_2 = tokenizer.encode_plus(sentence_0, sentence_2, add_special_tokens=True, return_tensors='pt')

pred_1 = pytorch_model(**inputs_1)[0].argmax().item()
pred_2 = pytorch_model(**inputs_2)[0].argmax().item()
print('sentence_1 is', 'a paraphrase' if pred_1 else 'not a paraphrase', 'of sentence_0')
print('sentence_2 is', 'a paraphrase' if pred_2 else 'not a paraphrase', 'of sentence_0')
```

3. See the following error message:
```
2019-11-01 11:57:45.714479: W tensorflow/core/common_runtime/base_collective_executor.cc:216] BaseCollectiveExecutor::StartAbort Invalid argument: Incompatible shapes: [512,512,32,12] vs. [512,1023,32,12]
         [[{{node tfxl_net_for_sequence_classification/transformer/layer_._0/rel_attn/add_3}}]]
  1/114 [..............................] - ETA: 1:02:58Traceback (most recent call last):
  File "run_tf_glue_xlnet.py", line 63, in <module>
    validation_data=valid_dataset, validation_steps=valid_steps)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training.py", line 728, in fit
    use_multiprocessing=use_multiprocessing)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 324, in fit
    total_epochs=epochs)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 123, in run_one_epoch
    batch_outs = execution_function(iterator)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_v2_utils.py", line 86, in execution_function
    distributed_function(input_fn))
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/def_function.py", line 457, in __call__
    result = self._call(*args, **kwds)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/def_function.py", line 520, in _call
    return self._stateless_fn(*args, **kwds)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 1823, in __call__
    return graph_function._filtered_call(args, kwargs)  # pylint: disable=protected-access
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 1141, in _filtered_call
    self.captured_inputs)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 1224, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 511, in call
    ctx=ctx)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/execute.py", line 67, in quick_execute
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError:  Incompatible shapes: [512,512,32,12] vs. [512,1023,32,12]
         [[node tfxl_net_for_sequence_classification/transformer/layer_._0/rel_attn/add_3 (defined at /data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py:1751) ]] [Op:__inference_distributed_function_67728]
```

## Expected behavior

The script should run just as `run_tf_glue.py`

## Environment

OS: CentOS-7
Python version: 3.6.9
PyTorch version: 1.2.0
PyTorch Transformers version (or branch): 2.1.1 (master from git)
Using GPU ? Yes
Distributed of parallel setup ? No
Any other relevant information:

## Additional context

<!-- Add any other context about the problem here. -->


Comments:
I have a fix for this in #1763 

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.