Head masking and test_head_masking not working properly for TFT5 models.
Author: stancld
Pull request: no

When removing `test_head_masking` flags during #9858, I found out `test_headmasking` was actually never run for `TFT5Model` and it seems there must be a bug, please see below:
```
_______________________________________________________________________________________________________ TFT5ModelTest.test_headmasking _______________________________________________________________________________________________________

self = <tests.test_modeling_tf_t5.TFT5ModelTest testMethod=test_headmasking>

    def test_headmasking(self):
        if not self.test_head_masking:
            return
    
        random.Random().seed(42)
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        random.Random().seed()
    
        inputs_dict["output_attentions"] = True
        config.output_hidden_states = True
        configs_no_init = _config_zero_init(config)  # To be sure we have no Nan
        for model_class in self.all_model_classes:
            model = model_class(config=configs_no_init)
    
            # Prepare head_mask
            def prepare_layer_head_mask(i, attention_heads, num_hidden_layers):
                if i == 0:
                    return tf.concat(
                        (tf.zeros(1, dtype=tf.float32), tf.ones(attention_heads - 1, dtype=tf.float32)), 0
                    )
                elif i == num_hidden_layers - 1:
                    return tf.concat(
                        (tf.zeros(attention_heads - 1, dtype=tf.float32), tf.ones(1, dtype=tf.float32)), 0
                    )
                else:
                    return tf.ones(attention_heads, dtype=tf.float32)
    
            head_mask = tf.stack(
                [
                    prepare_layer_head_mask(i, config.num_attention_heads, config.num_hidden_layers)
                    for i in range(config.num_hidden_layers)
                ],
                0,
            )
    
            inputs = self._prepare_for_class(inputs_dict, model_class).copy()
            inputs["head_mask"] = head_mask
            if model.config.is_encoder_decoder:
                signature = inspect.signature(model.call)
                arg_names = [*signature.parameters.keys()]
                if "decoder_head_mask" in arg_names:  # necessary diferentiation because of T5 model
                    inputs["decoder_head_mask"] = head_mask
    
>           outputs = model(**inputs, return_dict=True)

test_modeling_tf_common.py:686: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
../../../../../../miniconda3/envs/bart/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py:1012: in __call__
    outputs = call_fn(inputs, *args, **kwargs)
../src/transformers/models/t5/modeling_tf_t5.py:1160: in call
    inputs["encoder_outputs"] = self.encoder(
../../../../../../miniconda3/envs/bart/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py:1012: in __call__
    outputs = call_fn(inputs, *args, **kwargs)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <transformers.models.t5.modeling_tf_t5.TFT5MainLayer object at 0x7f8c38206a30>
input_ids = <tf.Tensor: shape=(13, 7), dtype=int32, numpy=
array([[63, 79, 60,  1, 57, 50, 42],
       [27,  6, 27, 88, 79, 14,  3... [95, 95, 79, 95, 63, 32, 24],
       [ 8,  9, 14, 46, 91, 75, 56],
       [26, 78, 52, 95, 45, 33, 78]], dtype=int32)>
attention_mask = None, encoder_hidden_states = None, encoder_attention_mask = None, inputs_embeds = None
head_mask = <tf.Tensor: shape=(5, 4), dtype=float32, numpy=
array([[0., 1., 1., 1.],
       [1., 1., 1., 1.],
       [1., 1., 1., 1.],
       [1., 1., 1., 1.],
       [0., 0., 0., 1.]], dtype=float32)>, encoder_head_mask = None
past_key_values = None, use_cache = False, output_attentions = True, output_hidden_states = True, return_dict = True, training = False, kwargs = {}
inputs = {'attention_mask': <tf.Tensor: shape=(13, 7), dtype=float32, numpy=
array([[1., 1., 1., 1., 1., 1., 1.],
       [1., 1..., 1.]], dtype=float32)>, 'encoder_attention_mask': None, 'encoder_head_mask': None, 'encoder_hidden_states': None, ...}
input_shape = [13, 7], batch_size = 13, seq_length = 7, mask_seq_length = 7

    def call(
        self,
        input_ids=None,
        attention_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        inputs_embeds=None,
        head_mask=None,
        encoder_head_mask=None,
        past_key_values=None,
        use_cache=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        training=False,
        **kwargs,
    ) -> Tuple:
        inputs = input_processing(
            func=self.call,
            config=self.config,
            input_ids=input_ids,
            attention_mask=attention_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            inputs_embeds=inputs_embeds,
            head_mask=head_mask,
            encoder_head_mask=encoder_head_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
            kwargs_call=kwargs,
        )
    
        if inputs["input_ids"] is not None and inputs["inputs_embeds"] is not None:
            err_msg_prefix = "decoder_" if self.is_decoder else ""
            raise ValueError(
                f"You cannot specify both {err_msg_prefix}inputs and {err_msg_prefix}inputs_embeds at the same time"
            )
        elif inputs["input_ids"] is not None:
            input_shape = shape_list(inputs["input_ids"])
            inputs["input_ids"] = tf.reshape(inputs["input_ids"], (-1, input_shape[-1]))
        elif inputs["inputs_embeds"] is not None:
            input_shape = shape_list(inputs["inputs_embeds"])[:-1]
        else:
            err_msg_prefix = "decoder_" if self.is_decoder else ""
            raise ValueError(f"You have to specify either {err_msg_prefix}inputs or {err_msg_prefix}inputs_embeds")
    
        if inputs["inputs_embeds"] is None:
            assert self.embed_tokens is not None, "You have to intialize the model with valid token embeddings"
            inputs["inputs_embeds"] = self.embed_tokens(inputs["input_ids"])
    
        batch_size, seq_length = input_shape
    
        # required mask seq length can be calculated via length of past
        mask_seq_length = (
            shape_list(inputs["past_key_values"][0][0])[2] + seq_length
            if inputs["past_key_values"] is not None
            else seq_length
        )
    
        if inputs["attention_mask"] is None:
            inputs["attention_mask"] = tf.fill((batch_size, mask_seq_length), 1)
        if (
            self.is_decoder
            and inputs["encoder_attention_mask"] is None
            and inputs["encoder_hidden_states"] is not None
        ):
            encoder_seq_length = shape_list(inputs["encoder_hidden_states"])[1]
            inputs["encoder_attention_mask"] = tf.fill((batch_size, encoder_seq_length), 1)
    
        # initialize past_key_values with `None` if past does not exist
        if inputs["past_key_values"] is None:
            inputs["past_key_values"] = [None] * len(self.block)
    
        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        inputs["attention_mask"] = tf.cast(inputs["attention_mask"], dtype=tf.float32)
        num_dims_attention_mask = len(shape_list(inputs["attention_mask"]))
        if num_dims_attention_mask == 3:
            extended_attention_mask = inputs["attention_mask"][:, None, :, :]
        elif num_dims_attention_mask == 2:
            # Provided a padding mask of dimensions [batch_size, mask_seq_length]
            # - if the model is a decoder, apply a causal mask in addition to the padding mask
            # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, mask_seq_length, mask_seq_length]
            if self.is_decoder:
                seq_ids = tf.range(mask_seq_length)
                causal_mask = tf.less_equal(
                    tf.tile(seq_ids[None, None, :], (batch_size, mask_seq_length, 1)),
                    seq_ids[None, :, None],
                )
                causal_mask = tf.cast(causal_mask, dtype=tf.float32)
                extended_attention_mask = causal_mask[:, None, :, :] * inputs["attention_mask"][:, None, None, :]
                if inputs["past_key_values"][0] is not None:
                    extended_attention_mask = extended_attention_mask[:, :, -seq_length:, :]
            else:
                extended_attention_mask = inputs["attention_mask"][:, None, None, :]
    
        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
        # masked positions, this operation will create a tensor which is 0.0 for
        # positions we want to attend and  -1e9 for masked positions.
        # Since we are adding it to the raw scores before the softmax, this is
        # effectively the same as removing these entirely.
    
        # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
        # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow/transformer/transformer_layers.py#L270
        # extended_attention_mask = tf.math.equal(extended_attention_mask,
        #                                         tf.transpose(extended_attention_mask, perm=(-1, -2)))
    
        extended_attention_mask = (1.0 - extended_attention_mask) * -1e9
    
        if self.is_decoder and inputs["encoder_attention_mask"] is not None:
            # If a 2D ou 3D attention mask is provided for the cross-attention
            # we need to make broadcastable to [batch_size, num_heads, mask_seq_length, mask_seq_length]
            # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
            inputs["encoder_attention_mask"] = tf.cast(inputs["encoder_attention_mask"], dtype=tf.float32)
            num_dims_encoder_attention_mask = len(shape_list(inputs["encoder_attention_mask"]))
            if num_dims_encoder_attention_mask == 3:
                encoder_extended_attention_mask = inputs["encoder_attention_mask"][:, None, :, :]
            if num_dims_encoder_attention_mask == 2:
                encoder_extended_attention_mask = inputs["encoder_attention_mask"][:, None, None, :]
    
            # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
            # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow/transformer/transformer_layers.py#L270
            # encoder_extended_attention_mask = tf.math.equal(encoder_extended_attention_mask,
            #                                         tf.transpose(encoder_extended_attention_mask, perm=(-1, -2)))
    
            encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * -1e9
        else:
            encoder_extended_attention_mask = None
    
>       assert inputs["head_mask"] is None, "Head mask not supported"
E       AssertionError: Head mask not supported

../src/transformers/models/t5/modeling_tf_t5.py:714: AssertionError
============================================================================================================== warnings summary ==============================================================================================================
../../../../../../miniconda3/envs/bart/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py:22
  /Users/daniel.stancl/miniconda3/envs/bart/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py:22: DeprecationWarning: the imp module is deprecated in favour of importlib; see the module's documentation for alternative uses
    import imp

tests/test_modeling_tf_t5.py: 44 warnings
  /var/folders/vs/4jsdk4nx1ds2m48ltfk3nmdc0000gn/T/tmpc35hmpmg.py:8: DeprecationWarning: The 'warn' method is deprecated, use 'warning' instead
    ag__.converted_call(ag__.ld(warnings).warn, ("The 'warn' method is deprecated, use 'warning' instead", ag__.ld(DeprecationWarning), 2), None, fscope)

tests/test_modeling_tf_t5.py::TFT5EncoderOnlyModelTest::test_saved_model_creation
tests/test_modeling_tf_t5.py::TFT5EncoderOnlyModelTest::test_saved_model_creation
tests/test_modeling_tf_t5.py::TFT5EncoderOnlyModelTest::test_saved_model_creation
tests/test_modeling_tf_t5.py::TFT5EncoderOnlyModelTest::test_saved_model_creation
tests/test_modeling_tf_t5.py::TFT5EncoderOnlyModelTest::test_saved_model_creation
  /Users/daniel.stancl/Documents/PhD/Projects/test_transformers/transformers/src/transformers/modeling_tf_utils.py:293: DeprecationWarning: The 'warn' method is deprecated, use 'warning' instead
    tf_logger.warn(

tests/test_modeling_tf_t5.py::TFT5EncoderOnlyModelTest::test_saved_model_creation
tests/test_modeling_tf_t5.py::TFT5EncoderOnlyModelTest::test_saved_model_creation
tests/test_modeling_tf_t5.py::TFT5EncoderOnlyModelTest::test_saved_model_creation
tests/test_modeling_tf_t5.py::TFT5EncoderOnlyModelTest::test_saved_model_creation
tests/test_modeling_tf_t5.py::TFT5EncoderOnlyModelTest::test_saved_model_creation
  /Users/daniel.stancl/Documents/PhD/Projects/test_transformers/transformers/src/transformers/modeling_tf_utils.py:302: DeprecationWarning: The 'warn' method is deprecated, use 'warning' instead
    tf_logger.warn("The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.")

-- Docs: https://docs.pytest.org/en/stable/warnings.html
========================================================================================================== short test summary info ===========================================================================================================
FAILED test_modeling_tf_t5.py::TFT5ModelTest::test_headmasking - AssertionError: Head mask not supported
```

My contribution: I'm gonna try to take care of this tomorrow.

<hr>
Reviewer: @jplu 

Comments:
That's an assert we put in T5 because the head mask is not supported. Happy you take care of this!!!