Add head_mask/decoder_head_mask for TF BART models
Author: stancld
Pull request: yes

This PR adds `head_mask` and `decoder_head_mask` input arguments for TF BART-based models. The full list of models is as follows:
* **TFBART**
* **TFMBart**
* **TFBlenderbot**
* **TFBlenderbotSmall**
* **TFMarian**
* **TFPegasus**

This PR can be deemed as a TF counterpart to the PR #9569. 

<hr>

**Further information:**
* I've added `test_headmasking` functionality to `tests/test_modeling_tf_common.py`

* **_TODO_**: Add a test (as a part of `test_headmasking`) to verify that we can get a gradient back for importance score computation. I am not so familiar with TensorFlow, therefore, I am not fully sure with a TF equivalent to
```
outputs = model(**inputs, return_dict=True)

output = sum(t.sum() for t in outputs[0])
output = output.sum()
output.backward()
```

<hr>

Reviewer: @patrickvonplaten 


Comments:
@stancld, thanks so much for tackling this! I think it would be a great addition if we could add a `test_headmasking` method for TF as well. 
I think it might be better if we don't try to have the exact test as in PyTorch. For now it should be enough to just leave out all gradient-related statements, like `.backward()`, `requires_grad(...)` in TF. The attentions output should still be 0 accordingly. 

Hey @patrickvonplaten, I hope this PR is ready for review. There's newly implemented `test_headmasking` method which follows the method from PyTorch testing except for the gradient-related statements as you pointed above.

It seems all checks have passed after rebasing this PR.

Also, @jplu it would be great if you could take a quick look if this is all serving compatible (I don't see a reason why it wouldn't be though)

Just done further tests on your PR and the changes are not graph compliant and the following slow tests are failing:

- test_saved_model_creation
- test_saved_model_creation_extended

One of the reasons is what @sgugger raised.

@jplu @sgugger Thank you very much for your comments and suggested solution. I'll try to fix these issues and send a new commit!

Hi @jplu, could you, please, review the changes in the code I've done to say whether assertions are done more appropriately now? :)
I've been also struggling to run (on my local) those four slow tests you mentioned last time, but I'm gonna have a look at that at the weekend if we're afraid of not passing.

I confirm that the assertions are done more appropriately now! Those four tests are among the most important one for the TF code base (they are run in slow mode because unfortunately they take some time to be executed).

If you need some help to make them pass, I will be happy to.

> @jplu I removed `global_rng` and leave it as it was before changes. Hopefully, now this PR is ready for a final review

Are these tests finally pass? :
* test_saved_model_with_hidden_states_output
* test_saved_model_with_attentions_output
* test_saved_model_creation
* test_saved_model_creation_extended

If yes, I will approve the PR :)

@jplu I ran these 4 aforementioned tests for BART and all those tests passed.

Merging, thanks a lot for your efforts @stancld!!