Getting incompatible shapes when using global_attention_mask in TFLongformerModel
Author: sli0111
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.8.2
- Platform: Linux-5.4.104+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.10
- PyTorch version (GPU?): 1.9.0+cu102 (False)
- Tensorflow version (GPU?): 2.5.0 (False)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: no

### Who can help
@patrickvonplaten

## Information

I am using TFLongformerModel.

The problem arises when using:
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [x] my own task or dataset: (give details below)

## To reproduce
The task is to classify the relationship between two phrases given the document the phrase is in.  There are a total of 8 classes.  Separators are added between the phrases and document.  I would like to add global attention to phrase 1 and phrase 2.

For example:
```

 phrase 1 
 phrase 2 
 entire document 

```

Checked shape of inputs to ensure they were the same.
```
input ids
tf.Tensor([    0 35151 19026 ...     1     1     1], shape=(1024,), dtype=int32)

attention mask
tf.Tensor([1 1 1 ... 0 0 0], shape=(1024,), dtype=int32)

global attention mask
tf.Tensor([1 1 1 ... 0 0 0], shape=(1024,), dtype=int32)
```

Example code for creating the model:
```
def create_model(train_layers=False, lr=5e-5):
    # Input Layers
    input_ids = Input(shape=(max_length,), name='input_ids', dtype='int32')
    attention_mask = Input(shape=(max_length,), name='attention_mask', dtype='int32') 
    global_attention_mask = Input(shape=(max_length,), name='global_attention_mask', dtype='int32')  

    # Transformer Layer
    X = transformer_model(input_ids=input_ids,
                          attention_mask=attention_mask,
                          global_attention_mask=global_attention_mask)[0] # Gets the embeddings before the [CLS] pooling

    # Deep Neural Net
    X = GlobalAveragePooling1D()(X)
    X = Dense(512, activation='relu')(X)
    X = Dense(512, activation='relu')(X)
    X = Dense(512, activation='relu')(X)
    X = Dense(256, activation='relu')(X)
    X = Dense(128, activation='relu')(X)
    X = Dense(64, activation='relu')(X)
    X = Dense(8, activation='softmax')(X)
    model = Model(inputs=[input_ids, attention_mask, global_attention_mask], outputs = X)

    if train_layers == False:
        for layer in model.layers[:3]:
            layer.trainable = False
    elif train_layers == True:
        for layer in model.layers[:3]:
            layer.trainable = True

    opt = tf.keras.optimizers.Adam(learning_rate=lr)

    # Compile the model
    model.compile(optimizer=opt, 
                  loss='sparse_categorical_crossentropy', 
                  metrics = ['sparse_categorical_accuracy'])
    
    # Print model summary
    model.summary()
    
    return model
```

Example code for training:
```
model = create_model(train_layers=False, lr=5e-05)
output = model.fit(x=[X.input_ids, X.attention_mask, X.global_attention_mask, 
                  y=y_resample, 
                  batch_size = 1,
                  epochs = 5)
```

Error:
```I
InvalidArgumentError:  Incompatible shapes: [2,1024,12,513] vs. [2,1024,12,522]
[[node model_8/tf_longformer_model/longformer/encoder/layer_._0/attention/self/dropout_1/dropout/Mul_1 (defined at /usr/local/lib/python3.7/dist-packages/transformers/models/longformer/modeling_tf_longformer.py:823) ]] [Op:__inference_train_function_618763]

Function call stack:
train_function
```

It seems the last dimension in the error message is related to the number of non-zero values in the global_attention_mask.  In the example above, there are 9 tokens that have a value of 1 in the global_attention_mask and the rest is zero.  (522-513=9)

## Expected behavior

The model should start training.


Comments:
It appears that when the number of examples is a multiple of the batch_size, the training starts (i.e. batch_size=2 and num_examples=6).  This doesn't work for a batch_size=1.  

Unfortunately for my full dataset, I won't be able to run a batch_size=2 without OOM.

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.

Hey @sli0111,

I'm very sorry, but I won't find time in the near future to debug this :-/ Gently pinging @Rocketknight1 in case you have some bandwidth to check out longformer training in TF

Hi @sli0111, sorry for the delay but I'm taking a look at this now! Are you still encountering the issue, and have you discovered anything else about it?

Hi @Rocketknight1 I was able to workaround the issue by reducing the number of examples to an even number and running a batch_size=2.  I was not able to avoid the error when I had a batch_size=1.

That's interesting - I'll mark this one down for testing and let you know what I find.

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.